#!/usr/bin/env python3 from __future__ import annotations import asyncio import argparse import base64 import csv import difflib import errno import hashlib import html import importlib import io import json import mailbox import mimetypes import os import pickle import posixpath import platform import re import secrets import shlex import shutil import site import sqlite3 import subprocess import sys import tempfile import time import unicodedata import venv import zipfile import xml.etree.ElementTree as ET from collections import defaultdict from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from contextlib import contextmanager from datetime import date, datetime, timedelta, timezone from email import policy from email.parser import BytesParser from email.utils import getaddresses, parsedate_to_datetime from html.parser import HTMLParser from pathlib import Path from typing import Callable, Iterator from urllib import error as urllib_error from urllib import request as urllib_request from zoneinfo import ZoneInfo _UNLOADED_DEPENDENCY = object() ACTIVE_WORKSPACE_ROOT: Path | None = None ACTIVATED_PLUGIN_SITE_PACKAGES: set[str] = set() # Third-party parsing dependencies load on demand so ordinary commands do not # depend on every parser being importable up front. charset_normalizer = _UNLOADED_DEPENDENCY extract_msg = _UNLOADED_DEPENDENCY openpyxl = _UNLOADED_DEPENDENCY xlrd = _UNLOADED_DEPENDENCY pdfplumber = _UNLOADED_DEPENDENCY DocxDocument = _UNLOADED_DEPENDENCY rtf_to_text = _UNLOADED_DEPENDENCY PilImage = _UNLOADED_DEPENDENCY pypff = _UNLOADED_DEPENDENCY try: import fcntl except Exception: # pragma: no cover - platform-specific locking fcntl = None try: import msvcrt except Exception: # pragma: no cover - platform-specific locking msvcrt = None TOOL_VERSION = "1.1.17" SCHEMA_VERSION = 26 SESSION_SCHEMA_VERSION = 2 REQUIREMENTS_VERSION = "2026-04-21-phase11-document-deduplication" TEMPLATE_SOURCE = "skills/tool-template/tools.py" PINNED_RUNTIME_REQUIREMENTS = ( "pdfplumber==0.11.9", "python-docx==1.2.0", "openpyxl==3.1.5", "xlrd==2.0.1", "extract-msg==0.55.0", "libpff-python==20231205", "striprtf==0.0.26", "Pillow==10.3.0", "charset-normalizer==3.4.7", ) MANUAL_FIELD_LOCKS_COLUMN = "manual_field_locks_json" LEGACY_METADATA_LOCKS_COLUMN = "locked_metadata_fields_json" CHUNK_TARGET_CHARS = 3200 CHUNK_OVERLAP_CHARS = 250 CONVERSATION_PREVIEW_MAX_CHARS = 180000 DEFAULT_PAGE_SIZE = 10 MAX_PAGE_SIZE = 100 BROWSE_MODE_DOCUMENTS = "documents" BROWSE_MODE_CONVERSATIONS = "conversations" BROWSE_MODE_ENTITIES = "entities" DEFAULT_BROWSE_MODE = BROWSE_MODE_DOCUMENTS DEFAULT_DOCUMENT_DISPLAY_COLUMNS = ( "content_type", "title", "author", "date_created", "control_number", ) DEFAULT_DISPLAY_COLUMNS = DEFAULT_DOCUMENT_DISPLAY_COLUMNS DEFAULT_CONVERSATION_DISPLAY_COLUMNS = ( "conversation_type", "title", "participants", "last_activity", "document_count", ) DEFAULT_ENTITY_DISPLAY_COLUMNS = ( "entity_type", "label", "primary_email", "document_count", "roles", ) MAX_SCOPE_DATASETS = 999 MAX_FILTER_EXPRESSION_LENGTH = 8192 MAX_FILTER_IN_LIST_ITEMS = 999 MAX_SAVED_SCOPES = 256 DEFAULT_CHUNK_PAGE_SIZE = 50 MAX_CHUNK_PAGE_SIZE = 200 GET_DOC_SUMMARY_CHARS = 1200 MAX_GET_DOC_CHUNKS = 10 MAX_GET_DOC_TEXT_CHARS = 30000 DEFAULT_CHUNK_SEARCH_TOP_K = 12 MAX_CHUNK_SEARCH_TOP_K = 50 DEFAULT_CHUNK_SEARCH_PER_DOC_CAP = 3 MAX_CHUNK_SEARCH_PER_DOC_CAP = 10 MAX_CHUNK_SEARCH_TEXT_CHARS = 100000 DEFAULT_AGGREGATE_LIMIT = 20 MAX_AGGREGATE_LIMIT = 200 CONTROL_NUMBER_PREFIX = "DOC" CONTROL_NUMBER_BATCH_WIDTH = 3 CONTROL_NUMBER_FAMILY_WIDTH = 8 BENCHMARK_ENABLED = os.environ.get("RETRIEVER_BENCHMARK") == "1" BENCHMARK_EVENTS: list[dict[str, object]] = [] if BENCHMARK_ENABLED: BENCHMARK_EVENTS.append({"name": "module_import_start", "ts": time.perf_counter()}) def benchmark_mark(name: str, **fields: object) -> None: if not BENCHMARK_ENABLED: return event: dict[str, object] = {"name": name, "ts": time.perf_counter()} if fields: event.update(fields) BENCHMARK_EVENTS.append(event) def benchmark_payload(**fields: object) -> dict[str, object]: events: list[dict[str, object]] = [] deltas: list[dict[str, object]] = [] for event in BENCHMARK_EVENTS: events.append({key: value for key, value in event.items() if key != "ts"}) for previous, current in zip(BENCHMARK_EVENTS, BENCHMARK_EVENTS[1:]): deltas.append( { "from": previous["name"], "to": current["name"], "delta_ms": round((float(current["ts"]) - float(previous["ts"])) * 1000.0, 3), } ) payload: dict[str, object] = { "enabled": BENCHMARK_ENABLED, "events": events, "deltas": deltas, } if fields: payload.update(fields) return payload def benchmark_emit(**fields: object) -> None: if not BENCHMARK_ENABLED: return sys.stderr.write(json.dumps({"_bench": benchmark_payload(**fields)}) + "\n") sys.stderr.flush() CONTROL_NUMBER_ATTACHMENT_WIDTH = 3 EMU_PER_PIXEL = 9525 IMAGE_NATIVE_PREVIEW_FILE_TYPES = { "bmp", "gif", "jpeg", "jpg", "png", "tif", "tiff", "webp", } CURATED_TEXT_SOURCE_FILE_TYPES = { "bash", "c", "cfg", "conf", "cpp", "cs", "css", "go", "h", "hpp", "ini", "java", "js", "jsx", "kt", "less", "php", "properties", "ps1", "py", "rb", "rs", "scss", "sh", "sql", "swift", "toml", "ts", "tsx", "xml", "yaml", "yml", "zsh", } SUPPORTED_FILE_TYPES = { "bash", "bmp", "c", "cfg", "conf", "cpp", "csv", "cs", "css", "docx", "eml", "gif", "go", "h", "htm", "html", "hpp", "ics", "ini", "java", "jpeg", "jpg", "json", "js", "jsx", "kt", "less", "mbox", "md", "msg", "pdf", "php", "pst", "png", "pptx", "properties", "ps1", "py", "rb", "rs", "rtf", "scss", "sh", "sql", "swift", "tif", "tiff", "toml", "tsv", "txt", "ts", "tsx", "webp", "xls", "xml", "xlsx", "yaml", "yml", "zsh", } NATIVE_PREVIEW_FILE_TYPES = { "bash", "bmp", "c", "cfg", "conf", "cpp", "csv", "cs", "css", "docx", "gif", "go", "h", "htm", "html", "hpp", "ics", "ini", "java", "jpeg", "jpg", "json", "js", "jsx", "kt", "less", "md", "pdf", "php", "png", "properties", "ps1", "py", "rb", "rs", "scss", "sh", "sql", "swift", "tif", "tiff", "toml", "tsv", "txt", "ts", "tsx", "webp", "xml", "yaml", "yml", "zsh", } TEXT_FILE_TYPES = {"csv", "htm", "html", "ics", "json", "md", "tsv", "txt", *CURATED_TEXT_SOURCE_FILE_TYPES} EDITABLE_BUILTIN_FIELDS = { "author", "content_type", "date_created", "date_modified", "page_count", "participants", "recipients", "subject", "title", } CONVERSATION_ASSIGNMENT_MODE_AUTO = "auto" CONVERSATION_ASSIGNMENT_MODE_MANUAL = "manual" SYSTEM_MANAGED_FIELDS = { "active_search_text_revision_id", "active_text_language", "active_text_quality_score", "active_text_source_kind", "canonical_kind", "canonical_status", "content_hash", "control_number_attachment_sequence", "control_number_batch", "control_number_family_sequence", "control_number", "conversation_id", "conversation_assignment_mode", "dataset_id", "file_hash", "file_name", "file_size", "file_type", "id", "ingested_at", "last_seen_at", "lifecycle_status", MANUAL_FIELD_LOCKS_COLUMN, LEGACY_METADATA_LOCKS_COLUMN, "merged_into_document_id", "begin_attachment", "begin_bates", "end_attachment", "end_bates", "child_document_kind", "parent_document_id", "production_id", "rel_path", "root_message_key", "source_folder_path", "source_item_id", "source_rel_path", "source_kind", "source_text_revision_id", "text_status", "updated_at", } BUILTIN_FIELD_TYPES = { "id": "integer", "canonical_kind": "text", "canonical_status": "text", "control_number": "text", "conversation_id": "integer", "conversation_assignment_mode": "text", "dataset_id": "integer", "merged_into_document_id": "integer", "parent_document_id": "integer", "child_document_kind": "text", "source_kind": "text", "source_rel_path": "text", "source_item_id": "text", "root_message_key": "text", "source_folder_path": "text", "production_id": "integer", "begin_bates": "text", "end_bates": "text", "begin_attachment": "text", "end_attachment": "text", "rel_path": "text", "file_name": "text", "file_type": "text", "file_size": "integer", "page_count": "integer", "author": "text", "content_type": "text", "date_created": "date", "date_modified": "date", "title": "text", "subject": "text", "participants": "text", "recipients": "text", MANUAL_FIELD_LOCKS_COLUMN: "text", "file_hash": "text", "content_hash": "text", "source_text_revision_id": "integer", "active_search_text_revision_id": "integer", "active_text_source_kind": "text", "active_text_language": "text", "active_text_quality_score": "real", "text_status": "text", "lifecycle_status": "text", "ingested_at": "date", "last_seen_at": "date", "updated_at": "date", "control_number_batch": "integer", "control_number_family_sequence": "integer", "control_number_attachment_sequence": "integer", } FIELD_NAME_ALIASES = { "dataset": "dataset_name", "dataset_label": "dataset_name", "collected_from": "custodian", "created_date": "date_created", "modified_date": "date_modified", } PASSIVE_FIELD_LABELS = { "active_search_text_revision_id": "Active Search Revision ID", "active_text_language": "Active Text Language", "active_text_quality_score": "Active Text Quality Score", "active_text_source_kind": "Active Text Source", "author": "Author", "begin_attachment": "Begin Attachment", "begin_bates": "Begin Bates", "canonical_kind": "Family", "canonical_status": "Record Status", "child_document_kind": "Document Role", "content_hash": "Content Hash", "content_type": "Type", "control_number": "Control #", "control_number_attachment_sequence": "Attachment Seq.", "control_number_batch": "Batch", "control_number_family_sequence": "Family Seq.", "conversation_assignment_mode": "Assignment Mode", "conversation_id": "Conversation ID", "conversation_type": "Type", "custodian": "Custodian", "dataset_id": "Dataset ID", "dataset_name": "Dataset", "date_created": "Created", "date_modified": "Modified", "document_count": "Documents", "end_attachment": "End Attachment", "end_bates": "End Bates", "file_hash": "File Hash", "file_name": "File", "file_size": "Size", "file_type": "File Type", "first_activity": "Started", "has_attachments": "Has Attachments", "id": "ID", "ingested_at": "Ingested", "is_attachment": "Attachment", "last_activity": "Last Activity", "last_seen_at": "Last Seen", "lifecycle_status": "File Status", "matching_document_count": "Matches", "merged_into_document_id": "Merged Into ID", "page_count": "Pages", "parent_document_id": "Parent ID", "participants": "Participants", "production_id": "Production ID", "production_name": "Production", "recipients": "Recipients", "rel_path": "Path", "root_message_key": "Root Message Key", "source_folder_path": "Source Folder", "source_item_id": "Source Item ID", "source_kind": "Source", "source_rel_path": "Source Path", "source_text_revision_id": "Source Text Revision ID", "subject": "Subject", "text_status": "Text Status", "title": "Title", "updated_at": "Updated", } PASSIVE_MIXED_CONTEXT_FIELD_LABELS = { "content_type": "Document Type", "conversation_type": "Conversation Type", } PASSIVE_FIELD_LABEL_UPPERCASE_TOKENS = { "api", "csv", "eml", "fts", "html", "id", "json", "mbox", "msg", "ocr", "pdf", "pst", "sql", "tsv", "uri", "url", "utc", "xml", } def passive_field_label(field_name: object, *, mixed_context: bool = False) -> str: normalized_name = str(field_name or "").strip() if not normalized_name: return "" canonical_name = FIELD_NAME_ALIASES.get(normalized_name, normalized_name) if mixed_context: mixed_context_label = PASSIVE_MIXED_CONTEXT_FIELD_LABELS.get(canonical_name) if mixed_context_label: return mixed_context_label label = PASSIVE_FIELD_LABELS.get(canonical_name) if label: return label words: list[str] = [] for token in re.split(r"[_\s]+", canonical_name): if not token: continue lower_token = token.lower() if lower_token in PASSIVE_FIELD_LABEL_UPPERCASE_TOKENS: words.append(lower_token.upper()) elif token.isupper(): words.append(token) else: words.append(token[:1].upper() + token[1:]) return " ".join(words) REGISTRY_FIELD_TYPES = { "boolean": "INTEGER", "date": "TEXT", "integer": "INTEGER", "real": "REAL", "text": "TEXT", } VIRTUAL_FILTER_FIELD_TYPES = { "custodian": "text", "dataset_name": "text", "is_attachment": "boolean", "has_attachments": "boolean", "participant": "text", "recipient": "text", "raw_author": "text", "raw_custodian": "text", "raw_participant": "text", "raw_participants": "text", "raw_recipient": "text", "raw_recipients": "text", "author_entity_id": "integer", "custodian_entity_id": "integer", "participant_entity_id": "integer", "participants_entity_id": "integer", "recipient_entity_id": "integer", "recipients_entity_id": "integer", "production_name": "text", } DISPLAYABLE_VIRTUAL_FIELDS = { "custodian", "dataset_name", "is_attachment", "production_name", } CATALOG_EXCLUDED_BUILTIN_FIELDS = { "active_search_text_revision_id", "active_text_language", "active_text_quality_score", "active_text_source_kind", "content_hash", "conversation_assignment_mode", "dataset_id", "file_hash", "id", MANUAL_FIELD_LOCKS_COLUMN, LEGACY_METADATA_LOCKS_COLUMN, "production_id", "root_message_key", "source_text_revision_id", } CATALOG_EXCLUDED_CUSTOM_FIELDS = { MANUAL_FIELD_LOCKS_COLUMN, LEGACY_METADATA_LOCKS_COLUMN, } AGGREGATABLE_VIRTUAL_FIELDS = {"dataset_name"} BUILTIN_FIELD_DESCRIPTIONS = { "author": "Document author from file metadata", "begin_bates": "Beginning Bates label for a production document", "canonical_kind": "Normalized logical content family used for dedupe safety checks", "canonical_status": "Logical document lifecycle state such as active, derelict, or merged", "content_type": "Normalized content category such as Email, E-Doc, or Chat", "control_number": "Stable document label used for review and export", "conversation_id": "Internal conversation grouping id shared by related documents", "conversation_assignment_mode": "Whether conversation grouping is automatic or manually pinned", "custodian": "Custodian or mailbox owner associated with the document", "date_created": "ISO date the document was created", "date_modified": "ISO date the document was last modified", "end_bates": "Ending Bates label for a production document", "file_name": "Document file name", "file_size": "File size in bytes", "file_type": "Normalized file extension such as pdf, docx, or eml", "merged_into_document_id": "Canonical survivor id when this document has been merged", "page_count": "Page or sheet count when available", "child_document_kind": "Contained-child semantics such as attachment or reply_thread", "participants": "Participants extracted from chat or email-style content", "recipients": "Recipients extracted from message metadata", "rel_path": "Workspace-relative document path", "source_folder_path": "Container-relative source folder path for container-derived items", "source_kind": "Origin kind such as filesystem, production, email_attachment, pst, or mbox", "source_rel_path": "Source-relative path used to derive the document", "subject": "Email or message subject line", "title": "Document title from metadata or content", "updated_at": "ISO timestamp when Retriever last updated the document row", } VIRTUAL_FIELD_DESCRIPTIONS = { "custodian": "Custodian or mailbox owner across the document's active collected copies", "dataset_name": "Logical dataset label for PSTs, MBOXes, productions, and manual document sets", "has_attachments": "Whether the document has one or more attachment child documents", "is_attachment": "Whether the document is an attachment child document", "production_name": "Friendly production name for production-derived documents", } INTERNAL_DOCUMENT_COLUMNS = { "custodians_json", } CONTENT_TYPE_EXTENSION_GROUPS = [ ("Email", "dbx eml emlx mbox msg nsf ost p7m p7s pst tnef vcf"), ( "Spreadsheet / Table", "123 csv dat dif gsheet numbers ods ots qpw slk sxc tsv uxdc wk1 wk3 wk4 wks wq1 xla xlam xlm xls xlsb xlsm xlsx xlt xltm xltx xlw", ), ("Presentation", "gslides key odp pez pot potm potx pps ppsm ppsx ppt pptm pptx prz sdd show shw sldm sldx sti"), ("Image", "bmp cut dds emf emz exf exif fax gif hpg hpgl ico iff jng jpeg jpg koala lbm pbm pcd pcx pgm plo plt png ppm prn psd ras sgi snp svg svgz targa tga tif tiff wbmp wdp webp wmf wmz xbm"), ("Audio", "3ga aac ac3 aif aifc aiff amr ape au awb dss dvf flac gsm m3u m4a m4p m4r mid midi mmf mp3 msv ogg opus pcm ra ram raw voc wav wma wv"), ("Video", "3g2 3gp amv asf avi divx drc flv gifv m2ts m2v m4v mkv mng mov mp4 mpeg mpg mxf nsv ogv qt rm rmvb roq ts vob webm wmv yuv"), ("Database", "accdb bak bson dbf dmp exp frm json ldf mdb mdf myd myi ndf odb ora pdb rdb sql"), ("Web", "htm html mht mhtml xhtm xhtml xml"), ("E-Doc", "doc docm docx dot dotm dotx eps fm one pdf ps pub rtf txt vdx vsd vsdm vsdx vss vst vsw vsx wpd wps"), ( "Source Code", "asm bash c cfg class coffee conf cpp cs css dart ear egg egg-info elm ex exs f90 fs go gradle groovy h hs ini ipynb jar java javafx jl jmod jnlp js jsh jsp jspx jsx jws kt less lisp lua m ml mm pas perl pex php pipfile pl pom properties ps1 pth py pyc pyd pyi pyo pyw pyz pyzw r rb requirementstxt rs sass scala scss sh sol swift toml ts tsx vb vue war whl yaml yml zsh", ), ("Container", "7z alzip bz2 cab e01 ex01 gz l01 lx01 rar tar z zip"), ("Calendar", "calendar ical icalendar ics ifb invite vcal vcs"), ("CAD", "3dxml asmdot drwdot dwg dxf easm easmx edrw edrwx eprt eprtx prtdot sldasm slddrw sldprt stl"), ("Message", "rsmf"), ] OOXML_RELATIONSHIP_NS = {"rels": "http://schemas.openxmlformats.org/package/2006/relationships"} PPTX_NAMESPACES = { "a": "http://schemas.openxmlformats.org/drawingml/2006/main", "cp": "http://schemas.openxmlformats.org/package/2006/metadata/core-properties", "dc": "http://purl.org/dc/elements/1.1/", "dcterms": "http://purl.org/dc/terms/", "p": "http://schemas.openxmlformats.org/presentationml/2006/main", "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships", } PPTX_NOTES_RELATIONSHIP_TYPE = "http://schemas.openxmlformats.org/officeDocument/2006/relationships/notesSlide" PRODUCTION_SOURCE_KIND = "production" EMAIL_ATTACHMENT_SOURCE_KIND = "email_attachment" EMAIL_CONVERSATION_SOURCE_KIND = "email_conversation" FILESYSTEM_SOURCE_KIND = "filesystem" MBOX_SOURCE_KIND = "mbox" PST_SOURCE_KIND = "pst" SLACK_EXPORT_SOURCE_KIND = "slack_export" MANUAL_DATASET_SOURCE_KIND = "manual" ACTIVE_OCCURRENCE_STATUS = "active" OCCURRENCE_LIFECYCLE_STATUSES = {"active", "superseded", "missing", "deleted"} CANONICAL_STATUS_ACTIVE = "active" CANONICAL_STATUS_DERELICT = "derelict" CANONICAL_STATUS_MERGED = "merged" ENTITY_TYPE_PERSON = "person" ENTITY_TYPE_ORGANIZATION = "organization" ENTITY_TYPE_SHARED_MAILBOX = "shared_mailbox" ENTITY_TYPE_SYSTEM_MAILBOX = "system_mailbox" ENTITY_TYPE_UNKNOWN = "unknown" ENTITY_TYPES = { ENTITY_TYPE_PERSON, ENTITY_TYPE_ORGANIZATION, ENTITY_TYPE_SHARED_MAILBOX, ENTITY_TYPE_SYSTEM_MAILBOX, ENTITY_TYPE_UNKNOWN, } SYSTEM_MAILBOX_LOCAL_PARTS = { "automated", "daemon", "donotreply", "mailerdaemon", "notification", "notifications", "noreply", "postmaster", } SYSTEM_MAILBOX_LOCAL_CONTAINS = { "donotreply", "noreply", } SHARED_MAILBOX_LOCAL_PARTS = { "admin", "allcompany", "billing", "contact", "events", "help", "hello", "info", "legal", "legalevents", "marketing", "news", "press", "sales", "support", "team", } SHARED_MAILBOX_NAME_HINTS = { "all company", "help desk", "legalweek", "support", "team", } ENTITY_ORIGIN_OBSERVED = "observed" ENTITY_ORIGIN_IDENTIFIED = "identified" ENTITY_ORIGIN_MANUAL = "manual" ENTITY_ORIGINS = { ENTITY_ORIGIN_OBSERVED, ENTITY_ORIGIN_IDENTIFIED, ENTITY_ORIGIN_MANUAL, } ENTITY_STATUS_ACTIVE = "active" ENTITY_STATUS_MERGED = "merged" ENTITY_STATUS_IGNORED = "ignored" ENTITY_STATUSES = { ENTITY_STATUS_ACTIVE, ENTITY_STATUS_MERGED, ENTITY_STATUS_IGNORED, } ENTITY_DISPLAY_SOURCE_AUTO = "auto" ENTITY_DISPLAY_SOURCE_MANUAL = "manual" ENTITY_DISPLAY_SOURCES = { ENTITY_DISPLAY_SOURCE_AUTO, ENTITY_DISPLAY_SOURCE_MANUAL, } ENTITY_IDENTIFIER_TYPES = { "email", "phone", "name", "handle", "external_id", } DOCUMENT_ENTITY_ROLES = { "author", "participant", "recipient", "custodian", } CANONICAL_KIND_VALUES = { "email", "document", "spreadsheet", "presentation", "image", "code", "data", "binary", "unknown", } TEXT_STATUS_PRIORITIES = { "ok": 0, "partial": 1, "empty": 2, "failed": 3, "error": 4, } SOURCE_KIND_PREFERRED_ORDER = { PRODUCTION_SOURCE_KIND: 0, FILESYSTEM_SOURCE_KIND: 1, PST_SOURCE_KIND: 2, MBOX_SOURCE_KIND: 3, SLACK_EXPORT_SOURCE_KIND: 4, EMAIL_ATTACHMENT_SOURCE_KIND: 5, } CHILD_DOCUMENT_KIND_ATTACHMENT = "attachment" CHILD_DOCUMENT_KIND_REPLY_THREAD = "reply_thread" ALLOWED_CHILD_DOCUMENT_KINDS = { CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_REPLY_THREAD, } PRODUCTION_DAT_HEADER_ALIASES = { "begbates": "begin_bates", "beginbates": "begin_bates", "begin bates": "begin_bates", "endbates": "end_bates", "end bates": "end_bates", "begattach": "begin_attachment", "beginattach": "begin_attachment", "begin attachment": "begin_attachment", "endattach": "end_attachment", "endattachment": "end_attachment", "end attachment": "end_attachment", "filepath": "native_path", "file path": "native_path", "nativefile": "native_path", "native file": "native_path", "native path": "native_path", "textpath": "text_path", "text path": "text_path", "textprecedence": "text_path", "text precedence": "text_path", } def build_content_type_by_extension() -> dict[str, str]: mapping: dict[str, str] = {} for content_type, raw_extensions in CONTENT_TYPE_EXTENSION_GROUPS: for extension in raw_extensions.split(): normalized_extension = extension.strip().lower().strip(",") if normalized_extension: mapping.setdefault(normalized_extension, content_type) return mapping CONTENT_TYPE_BY_EXTENSION = build_content_type_by_extension() ATTACHMENT_SUFFIX_PATTERN = re.compile( r"^(?P.+?)\s+(?P<label>(?:and\s+(?:back\s*up|backup)\s+)?attachments?)\s*:\s*(?P<attachments>.+)$", re.IGNORECASE, ) HTML_PREVIEW_ATTACHMENT_LINKS_PATTERN = re.compile( r"<!-- RETRIEVER_ATTACHMENT_LINKS_START -->.*?<!-- RETRIEVER_ATTACHMENT_LINKS_END -->", re.DOTALL, ) HTML_PREVIEW_CALENDAR_INVITES_PATTERN = re.compile( r"<!-- RETRIEVER_CALENDAR_INVITES_START -->.*?<!-- RETRIEVER_CALENDAR_INVITES_END -->", re.DOTALL, ) SCHEMA_STATEMENTS = [ """ CREATE TABLE IF NOT EXISTS workspace_meta ( id INTEGER PRIMARY KEY CHECK (id = 1), schema_version INTEGER NOT NULL, tool_version TEXT NOT NULL, requirements_version TEXT NOT NULL, template_source TEXT NOT NULL, template_sha256 TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS datasets ( id INTEGER PRIMARY KEY, source_kind TEXT NOT NULL, dataset_locator TEXT NOT NULL, dataset_name TEXT NOT NULL, dataset_name_normalized TEXT, allow_auto_merge INTEGER NOT NULL DEFAULT 1, email_auto_merge INTEGER NOT NULL DEFAULT 1, handle_auto_merge INTEGER NOT NULL DEFAULT 1, phone_auto_merge INTEGER NOT NULL DEFAULT 0, name_auto_merge INTEGER NOT NULL DEFAULT 0, external_id_auto_merge_names_json TEXT NOT NULL DEFAULT '[]', created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS dataset_sources ( id INTEGER PRIMARY KEY, dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE, source_kind TEXT NOT NULL, source_locator TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS dataset_documents ( id INTEGER PRIMARY KEY, dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, dataset_source_id INTEGER REFERENCES dataset_sources(id) ON DELETE CASCADE, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS conversations ( id INTEGER PRIMARY KEY, source_kind TEXT NOT NULL, source_locator TEXT NOT NULL, conversation_key TEXT NOT NULL, conversation_type TEXT NOT NULL, display_name TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS documents ( id INTEGER PRIMARY KEY, control_number TEXT UNIQUE, canonical_kind TEXT NOT NULL DEFAULT 'unknown', canonical_status TEXT NOT NULL DEFAULT 'active', merged_into_document_id INTEGER REFERENCES documents(id) ON DELETE SET NULL, conversation_id INTEGER REFERENCES conversations(id) ON DELETE SET NULL, conversation_assignment_mode TEXT NOT NULL DEFAULT 'auto', dataset_id INTEGER REFERENCES datasets(id) ON DELETE SET NULL, parent_document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE, child_document_kind TEXT, source_kind TEXT, source_rel_path TEXT, source_item_id TEXT, root_message_key TEXT, source_folder_path TEXT, production_id INTEGER REFERENCES productions(id) ON DELETE SET NULL, begin_bates TEXT, end_bates TEXT, begin_attachment TEXT, end_attachment TEXT, rel_path TEXT NOT NULL UNIQUE, file_name TEXT NOT NULL, file_type TEXT, file_size INTEGER, page_count INTEGER, author TEXT, content_type TEXT, custodians_json TEXT NOT NULL DEFAULT '[]', date_created TEXT, date_modified TEXT, title TEXT, subject TEXT, participants TEXT, recipients TEXT, manual_field_locks_json TEXT NOT NULL DEFAULT '[]', file_hash TEXT, content_hash TEXT, text_status TEXT NOT NULL DEFAULT 'ok', lifecycle_status TEXT NOT NULL DEFAULT 'active', ingested_at TEXT, last_seen_at TEXT, updated_at TEXT, control_number_batch INTEGER, control_number_family_sequence INTEGER, control_number_attachment_sequence INTEGER ) """, """ CREATE TABLE IF NOT EXISTS document_occurrences ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, dataset_source_id INTEGER REFERENCES dataset_sources(id) ON DELETE SET NULL, parent_occurrence_id INTEGER REFERENCES document_occurrences(id) ON DELETE SET NULL, occurrence_control_number TEXT, source_kind TEXT, source_rel_path TEXT, source_item_id TEXT, source_folder_path TEXT, production_id INTEGER REFERENCES productions(id) ON DELETE SET NULL, begin_bates TEXT, end_bates TEXT, begin_attachment TEXT, end_attachment TEXT, rel_path TEXT NOT NULL, file_name TEXT NOT NULL, file_type TEXT, mime_type TEXT, file_size INTEGER, file_hash TEXT, custodian TEXT, fs_created_at TEXT, fs_modified_at TEXT, extracted_author TEXT, extracted_title TEXT, extracted_subject TEXT, extracted_participants TEXT, extracted_recipients TEXT, extracted_doc_authored_at TEXT, extracted_doc_modified_at TEXT, extracted_content_type TEXT, extracted_kind TEXT, entity_hints_json TEXT NOT NULL DEFAULT '{}', text_status TEXT NOT NULL DEFAULT 'ok', lifecycle_status TEXT NOT NULL DEFAULT 'active', has_preview INTEGER NOT NULL DEFAULT 0, ingested_at TEXT, last_seen_at TEXT, updated_at TEXT ) """, """ CREATE TABLE IF NOT EXISTS entities ( id INTEGER PRIMARY KEY, entity_type TEXT NOT NULL DEFAULT 'person', display_name TEXT, primary_email TEXT, primary_phone TEXT, sort_name TEXT, notes TEXT, display_name_source TEXT NOT NULL DEFAULT 'auto', entity_origin TEXT NOT NULL DEFAULT 'observed', canonical_status TEXT NOT NULL DEFAULT 'active', merged_into_entity_id INTEGER REFERENCES entities(id) ON DELETE SET NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, CHECK (entity_type IN ('person', 'organization', 'shared_mailbox', 'system_mailbox', 'unknown')), CHECK (display_name_source IN ('auto', 'manual')), CHECK (entity_origin IN ('observed', 'identified', 'manual')), CHECK (canonical_status IN ('active', 'merged', 'ignored')), CHECK ( (canonical_status = 'merged' AND merged_into_entity_id IS NOT NULL) OR (canonical_status != 'merged' AND merged_into_entity_id IS NULL) ) ) """, """ CREATE TABLE IF NOT EXISTS entity_identifiers ( id INTEGER PRIMARY KEY, entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE, identifier_type TEXT NOT NULL, display_value TEXT NOT NULL, normalized_value TEXT NOT NULL, provider TEXT, provider_scope TEXT, identifier_name TEXT, identifier_scope TEXT, parsed_name_json TEXT, parsed_phone_json TEXT, normalized_full_name TEXT, normalized_sort_name TEXT, is_primary INTEGER NOT NULL DEFAULT 0, is_verified INTEGER NOT NULL DEFAULT 0, source_kind TEXT NOT NULL DEFAULT 'auto', created_at TEXT NOT NULL, updated_at TEXT NOT NULL, CHECK (identifier_type IN ('email', 'phone', 'name', 'handle', 'external_id')) ) """, """ CREATE TABLE IF NOT EXISTS entity_resolution_keys ( id INTEGER PRIMARY KEY, entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE, identifier_id INTEGER REFERENCES entity_identifiers(id) ON DELETE CASCADE, key_type TEXT NOT NULL, provider TEXT, provider_scope TEXT, identifier_name TEXT, identifier_scope TEXT, normalized_value TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_entities ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE, role TEXT NOT NULL, ordinal INTEGER NOT NULL DEFAULT 0, assignment_mode TEXT NOT NULL DEFAULT 'auto', observed_title TEXT, evidence_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, updated_at TEXT NOT NULL, CHECK (role IN ('author', 'participant', 'recipient', 'custodian')), CHECK (assignment_mode IN ('auto', 'manual')) ) """, """ CREATE TABLE IF NOT EXISTS entity_overrides ( id INTEGER PRIMARY KEY, scope_type TEXT NOT NULL, scope_id INTEGER, role TEXT, source_entity_id INTEGER REFERENCES entities(id) ON DELETE SET NULL, normalized_candidate_key TEXT, replacement_entity_id INTEGER REFERENCES entities(id) ON DELETE SET NULL, override_effect TEXT NOT NULL, source_hint TEXT, reason TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, CHECK (scope_type IN ('document', 'global')), CHECK (override_effect IN ('replace', 'remove', 'ignore')), CHECK ( (scope_type = 'document' AND scope_id IS NOT NULL AND override_effect IN ('replace', 'remove')) OR (scope_type = 'global' AND scope_id IS NULL AND override_effect = 'ignore') ), CHECK ( (override_effect = 'replace' AND replacement_entity_id IS NOT NULL) OR (override_effect IN ('remove', 'ignore') AND replacement_entity_id IS NULL) ), CHECK ( override_effect != 'ignore' OR source_entity_id IS NOT NULL OR normalized_candidate_key IS NOT NULL ) ) """, """ CREATE TABLE IF NOT EXISTS entity_merge_blocks ( id INTEGER PRIMARY KEY, left_entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE, right_entity_id INTEGER NOT NULL REFERENCES entities(id) ON DELETE CASCADE, reason TEXT, created_at TEXT NOT NULL, CHECK (left_entity_id < right_entity_id) ) """, """ CREATE TABLE IF NOT EXISTS document_dedupe_keys ( id INTEGER PRIMARY KEY, basis TEXT NOT NULL, key_value TEXT NOT NULL, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS canonical_metadata_conflicts ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, field_name TEXT NOT NULL, occurrence_id INTEGER NOT NULL REFERENCES document_occurrences(id) ON DELETE CASCADE, value TEXT, first_seen_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_control_number_aliases ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, occurrence_id INTEGER REFERENCES document_occurrences(id) ON DELETE CASCADE, alias_value TEXT NOT NULL, alias_type TEXT NOT NULL, active_flag INTEGER NOT NULL DEFAULT 1, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_merge_events ( id INTEGER PRIMARY KEY, survivor_document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, loser_document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, merge_basis TEXT NOT NULL, actor TEXT, schema_version INTEGER NOT NULL, pre_merge_survivor_json TEXT NOT NULL, pre_merge_loser_json TEXT NOT NULL, artifact_counts_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_field_conflicts ( id INTEGER PRIMARY KEY, merge_event_id INTEGER REFERENCES document_merge_events(id) ON DELETE CASCADE, document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE, field_name TEXT NOT NULL, survivor_value TEXT, loser_value TEXT, resolution TEXT, created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_email_threading ( document_id INTEGER PRIMARY KEY REFERENCES documents(id) ON DELETE CASCADE, message_id TEXT, in_reply_to TEXT, references_json TEXT NOT NULL DEFAULT '[]', conversation_index TEXT, conversation_topic TEXT, normalized_subject TEXT, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_chat_threading ( document_id INTEGER PRIMARY KEY REFERENCES documents(id) ON DELETE CASCADE, thread_id TEXT, message_id TEXT, parent_message_id TEXT, thread_type TEXT, participants_json TEXT NOT NULL DEFAULT '[]', updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS productions ( id INTEGER PRIMARY KEY, dataset_id INTEGER REFERENCES datasets(id) ON DELETE SET NULL, rel_root TEXT NOT NULL UNIQUE, production_name TEXT NOT NULL, metadata_load_rel_path TEXT NOT NULL, image_load_rel_path TEXT, source_type TEXT NOT NULL, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_source_parts ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, part_kind TEXT NOT NULL, rel_source_path TEXT NOT NULL, ordinal INTEGER NOT NULL DEFAULT 0, label TEXT, created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS container_sources ( id INTEGER PRIMARY KEY, dataset_id INTEGER REFERENCES datasets(id) ON DELETE SET NULL, source_kind TEXT NOT NULL, source_rel_path TEXT NOT NULL UNIQUE, file_size INTEGER, file_mtime TEXT, file_hash TEXT, message_count INTEGER, last_scan_started_at TEXT, last_scan_completed_at TEXT, last_ingested_at TEXT ) """, """ CREATE TABLE IF NOT EXISTS document_previews ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, rel_preview_path TEXT NOT NULL, preview_type TEXT NOT NULL, target_fragment TEXT, label TEXT, ordinal INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS document_chunks ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, chunk_index INTEGER NOT NULL, char_start INTEGER NOT NULL, char_end INTEGER NOT NULL, token_estimate INTEGER, text_content TEXT NOT NULL, UNIQUE(document_id, chunk_index) ) """, """ CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5( document_id UNINDEXED, file_name, title, subject, author, custodian, participants, recipients ) """, """ CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5( chunk_id UNINDEXED, document_id UNINDEXED, text_content ) """, """ CREATE TABLE IF NOT EXISTS custom_fields_registry ( id INTEGER PRIMARY KEY, field_name TEXT NOT NULL UNIQUE, field_type TEXT NOT NULL, instruction TEXT, created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS jobs ( id INTEGER PRIMARY KEY, job_name TEXT NOT NULL UNIQUE, job_kind TEXT NOT NULL, description TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, archived_at TEXT ) """, """ CREATE TABLE IF NOT EXISTS job_outputs ( id INTEGER PRIMARY KEY, job_id INTEGER NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, output_name TEXT NOT NULL, value_type TEXT NOT NULL DEFAULT 'text', bound_custom_field TEXT, description TEXT, ordinal INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, UNIQUE(job_id, output_name) ) """, """ CREATE TABLE IF NOT EXISTS job_versions ( id INTEGER PRIMARY KEY, job_id INTEGER NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, version INTEGER NOT NULL, display_name TEXT NOT NULL, instruction_text TEXT NOT NULL DEFAULT '', instruction_hash TEXT NOT NULL, response_schema_json TEXT, capability TEXT NOT NULL, provider TEXT NOT NULL, model TEXT, parameters_json TEXT NOT NULL DEFAULT '{}', input_basis TEXT NOT NULL, segment_profile TEXT, aggregation_strategy TEXT, created_at TEXT NOT NULL, archived_at TEXT, UNIQUE(job_id, version) ) """, """ CREATE TABLE IF NOT EXISTS text_revisions ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, revision_kind TEXT NOT NULL, language TEXT, parent_revision_id INTEGER REFERENCES text_revisions(id) ON DELETE SET NULL, created_by_job_version_id INTEGER REFERENCES job_versions(id) ON DELETE SET NULL, storage_rel_path TEXT, content_hash TEXT NOT NULL, char_count INTEGER, token_estimate INTEGER, quality_score REAL, provider_metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, retracted_at TEXT, retraction_reason TEXT ) """, """ CREATE TABLE IF NOT EXISTS text_revision_segments ( id INTEGER PRIMARY KEY, revision_id INTEGER NOT NULL REFERENCES text_revisions(id) ON DELETE CASCADE, segment_profile TEXT NOT NULL, level INTEGER NOT NULL DEFAULT 0, parent_segment_id INTEGER REFERENCES text_revision_segments(id) ON DELETE CASCADE, ordinal INTEGER NOT NULL, char_start INTEGER NOT NULL, char_end INTEGER NOT NULL, token_estimate INTEGER, text_hash TEXT NOT NULL, created_at TEXT NOT NULL, UNIQUE(revision_id, segment_profile, level, ordinal) ) """, """ CREATE TABLE IF NOT EXISTS runs ( id INTEGER PRIMARY KEY, job_version_id INTEGER NOT NULL REFERENCES job_versions(id) ON DELETE CASCADE, from_run_id INTEGER REFERENCES runs(id) ON DELETE SET NULL, selector_json TEXT NOT NULL, exclude_selector_json TEXT NOT NULL DEFAULT '{}', activation_policy TEXT NOT NULL DEFAULT 'manual', family_mode TEXT NOT NULL DEFAULT 'exact', seed_limit INTEGER, status TEXT NOT NULL DEFAULT 'planned', planned_count INTEGER NOT NULL DEFAULT 0, completed_count INTEGER NOT NULL DEFAULT 0, failed_count INTEGER NOT NULL DEFAULT 0, skipped_count INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT, canceled_at TEXT ) """, """ CREATE TABLE IF NOT EXISTS run_snapshot_documents ( id INTEGER PRIMARY KEY, run_id INTEGER NOT NULL REFERENCES runs(id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, ordinal INTEGER NOT NULL, inclusion_reason_json TEXT NOT NULL DEFAULT '{}', pinned_input_revision_id INTEGER REFERENCES text_revisions(id) ON DELETE SET NULL, pinned_input_identity TEXT NOT NULL, pinned_content_hash TEXT, created_at TEXT NOT NULL, UNIQUE(run_id, document_id) ) """, """ CREATE TABLE IF NOT EXISTS run_items ( id INTEGER PRIMARY KEY, run_id INTEGER NOT NULL REFERENCES runs(id) ON DELETE CASCADE, run_snapshot_document_id INTEGER REFERENCES run_snapshot_documents(id) ON DELETE CASCADE, item_kind TEXT NOT NULL, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, page_number INTEGER, segment_id INTEGER REFERENCES text_revision_segments(id) ON DELETE CASCADE, input_artifact_rel_path TEXT, input_identity TEXT NOT NULL, result_id INTEGER REFERENCES results(id) ON DELETE SET NULL, status TEXT NOT NULL DEFAULT 'pending', claimed_by TEXT, claimed_at TEXT, last_heartbeat_at TEXT, attempt_count INTEGER NOT NULL DEFAULT 0, last_error TEXT, created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT ) """, """ CREATE TABLE IF NOT EXISTS run_workers ( id INTEGER PRIMARY KEY, run_id INTEGER NOT NULL REFERENCES runs(id) ON DELETE CASCADE, claimed_by TEXT NOT NULL, launch_mode TEXT NOT NULL DEFAULT 'inline', worker_task_id TEXT, status TEXT NOT NULL DEFAULT 'active', max_batches INTEGER, batches_prepared INTEGER NOT NULL DEFAULT 0, items_completed INTEGER NOT NULL DEFAULT 0, items_failed INTEGER NOT NULL DEFAULT 0, last_heartbeat_at TEXT, last_error TEXT, created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT, cancel_requested_at TEXT, summary_json TEXT NOT NULL DEFAULT '{}', UNIQUE(run_id, claimed_by) ) """, """ CREATE TABLE IF NOT EXISTS ocr_page_outputs ( id INTEGER PRIMARY KEY, run_item_id INTEGER NOT NULL REFERENCES run_items(id) ON DELETE CASCADE, run_id INTEGER NOT NULL REFERENCES runs(id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, page_number INTEGER NOT NULL, text_content TEXT NOT NULL, raw_output_json TEXT, normalized_output_json TEXT, provider_metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, UNIQUE(run_item_id) ) """, """ CREATE TABLE IF NOT EXISTS image_description_page_outputs ( id INTEGER PRIMARY KEY, run_item_id INTEGER NOT NULL REFERENCES run_items(id) ON DELETE CASCADE, run_id INTEGER NOT NULL REFERENCES runs(id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, page_number INTEGER NOT NULL, text_content TEXT NOT NULL, raw_output_json TEXT, normalized_output_json TEXT, provider_metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, UNIQUE(run_item_id) ) """, """ CREATE TABLE IF NOT EXISTS attempts ( id INTEGER PRIMARY KEY, run_item_id INTEGER NOT NULL REFERENCES run_items(id) ON DELETE CASCADE, attempt_number INTEGER NOT NULL, provider_request_id TEXT, input_tokens INTEGER, output_tokens INTEGER, cost_cents INTEGER, latency_ms INTEGER, provider_metadata_json TEXT NOT NULL DEFAULT '{}', error_summary TEXT, created_at TEXT NOT NULL, UNIQUE(run_item_id, attempt_number) ) """, """ CREATE TABLE IF NOT EXISTS results ( id INTEGER PRIMARY KEY, run_id INTEGER REFERENCES runs(id) ON DELETE SET NULL, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, job_version_id INTEGER NOT NULL REFERENCES job_versions(id) ON DELETE CASCADE, input_revision_id INTEGER REFERENCES text_revisions(id) ON DELETE SET NULL, input_identity TEXT NOT NULL, raw_output_json TEXT, normalized_output_json TEXT, created_text_revision_id INTEGER REFERENCES text_revisions(id) ON DELETE SET NULL, provider_metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, retracted_at TEXT, retraction_reason TEXT ) """, """ CREATE TABLE IF NOT EXISTS result_outputs ( id INTEGER PRIMARY KEY, result_id INTEGER NOT NULL REFERENCES results(id) ON DELETE CASCADE, job_output_id INTEGER NOT NULL REFERENCES job_outputs(id) ON DELETE CASCADE, output_value_json TEXT NOT NULL, display_value TEXT, score REAL, created_at TEXT NOT NULL, UNIQUE(result_id, job_output_id) ) """, """ CREATE TABLE IF NOT EXISTS embedding_vectors ( id INTEGER PRIMARY KEY, job_version_id INTEGER NOT NULL REFERENCES job_versions(id) ON DELETE CASCADE, revision_id INTEGER REFERENCES text_revisions(id) ON DELETE CASCADE, segment_id INTEGER NOT NULL REFERENCES text_revision_segments(id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, level INTEGER NOT NULL DEFAULT 0, vector_blob BLOB NOT NULL, encoding TEXT NOT NULL DEFAULT 'float32-le', dimensions INTEGER NOT NULL, distance_metric TEXT NOT NULL DEFAULT 'cosine', provider_metadata_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, retracted_at TEXT, retraction_reason TEXT ) """, """ CREATE TABLE IF NOT EXISTS publications ( id INTEGER PRIMARY KEY, result_output_id INTEGER NOT NULL REFERENCES result_outputs(id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, job_output_id INTEGER NOT NULL REFERENCES job_outputs(id) ON DELETE CASCADE, custom_field_name TEXT NOT NULL, published_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS text_revision_activation_events ( id INTEGER PRIMARY KEY, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, text_revision_id INTEGER NOT NULL REFERENCES text_revisions(id) ON DELETE CASCADE, activated_by_job_version_id INTEGER REFERENCES job_versions(id) ON DELETE SET NULL, source_result_id INTEGER REFERENCES results(id) ON DELETE SET NULL, activation_policy TEXT, created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS control_number_batches ( batch_number INTEGER PRIMARY KEY, next_family_sequence INTEGER NOT NULL DEFAULT 1, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS ingest_runs ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL UNIQUE, scope_json TEXT NOT NULL DEFAULT '{}', recursive INTEGER NOT NULL DEFAULT 0, raw_file_types TEXT, pipeline_schema_version INTEGER NOT NULL, phase TEXT NOT NULL DEFAULT 'planning', status TEXT NOT NULL DEFAULT 'planning', prepare_worker_soft_limit INTEGER NOT NULL DEFAULT 4, committer_lease_owner TEXT, committer_lease_expires_at TEXT, committer_heartbeat_at TEXT, entity_graph_stale INTEGER NOT NULL DEFAULT 0, entity_policy_snapshot_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT, cancel_requested_at TEXT, last_heartbeat_at TEXT, error TEXT ) """, """ CREATE TABLE IF NOT EXISTS ingest_work_items ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES ingest_runs(run_id) ON DELETE CASCADE, unit_type TEXT NOT NULL, source_kind TEXT, source_key TEXT, rel_path TEXT, commit_order INTEGER, parent_order INTEGER, spawned_by_work_item_id INTEGER REFERENCES ingest_work_items(id) ON DELETE SET NULL, payload_json TEXT NOT NULL DEFAULT '{}', affected_document_ids_json TEXT NOT NULL DEFAULT '[]', affected_conversation_keys_json TEXT NOT NULL DEFAULT '[]', affected_entity_ids_json TEXT NOT NULL DEFAULT '[]', artifact_manifest_json TEXT NOT NULL DEFAULT '{}', status TEXT NOT NULL DEFAULT 'pending', lease_owner TEXT, lease_expires_at TEXT, attempts INTEGER NOT NULL DEFAULT 0, last_error TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS ingest_prepared_items ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES ingest_runs(run_id) ON DELETE CASCADE, work_item_id INTEGER NOT NULL REFERENCES ingest_work_items(id) ON DELETE CASCADE, payload_kind TEXT NOT NULL, payload_json TEXT NOT NULL DEFAULT '{}', spill_rel_path TEXT, payload_bytes INTEGER NOT NULL DEFAULT 0, source_fingerprint_json TEXT NOT NULL DEFAULT '{}', prepared_at TEXT NOT NULL, error_json TEXT NOT NULL DEFAULT '{}', UNIQUE(work_item_id) ) """, """ CREATE TABLE IF NOT EXISTS ingest_rename_consumptions ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES ingest_runs(run_id) ON DELETE CASCADE, target_work_item_id INTEGER NOT NULL REFERENCES ingest_work_items(id) ON DELETE CASCADE, source_document_id INTEGER, source_occurrence_id INTEGER, file_hash TEXT NOT NULL, created_at TEXT NOT NULL, UNIQUE(run_id, source_occurrence_id) ) """, """ CREATE TABLE IF NOT EXISTS ingest_phase_cursors ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES ingest_runs(run_id) ON DELETE CASCADE, phase TEXT NOT NULL, cursor_key TEXT NOT NULL, cursor_json TEXT NOT NULL DEFAULT '{}', status TEXT NOT NULL DEFAULT 'pending', updated_at TEXT NOT NULL, UNIQUE(run_id, phase, cursor_key) ) """, """ CREATE TABLE IF NOT EXISTS ingest_worker_events ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES ingest_runs(run_id) ON DELETE CASCADE, worker_id TEXT, event_type TEXT NOT NULL, work_item_id INTEGER REFERENCES ingest_work_items(id) ON DELETE SET NULL, phase TEXT, duration_ms REAL, details_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS ingest_artifact_sweeps ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES ingest_runs(run_id) ON DELETE CASCADE, work_item_id INTEGER REFERENCES ingest_work_items(id) ON DELETE SET NULL, artifact_kind TEXT NOT NULL, temp_rel_path TEXT, canonical_rel_path TEXT, content_hash TEXT, state TEXT NOT NULL DEFAULT 'staged', created_at TEXT NOT NULL, updated_at TEXT NOT NULL ) """, """ CREATE TABLE IF NOT EXISTS entity_rebuild_runs ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL UNIQUE, mode TEXT NOT NULL DEFAULT 'full', phase TEXT NOT NULL DEFAULT 'resetting', status TEXT NOT NULL DEFAULT 'resetting', document_ids_json TEXT NOT NULL DEFAULT '[]', batch_size INTEGER NOT NULL DEFAULT 500, reset_stage TEXT NOT NULL DEFAULT 'document_entities', reset_counts_json TEXT NOT NULL DEFAULT '{}', cursor_json TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT, cancel_requested_at TEXT, last_heartbeat_at TEXT, error TEXT ) """, """ CREATE TABLE IF NOT EXISTS entity_rebuild_items ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES entity_rebuild_runs(run_id) ON DELETE CASCADE, document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE, ordinal INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL DEFAULT 'pending', lease_owner TEXT, lease_expires_at TEXT, attempts INTEGER NOT NULL DEFAULT 0, document_synced INTEGER NOT NULL DEFAULT 0, auto_links_created INTEGER NOT NULL DEFAULT 0, last_error TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, UNIQUE(run_id, document_id) ) """, """ CREATE TABLE IF NOT EXISTS export_runs ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL UNIQUE, export_kind TEXT NOT NULL, output_path TEXT NOT NULL, output_rel_path TEXT, selector_json TEXT NOT NULL DEFAULT '{}', config_json TEXT NOT NULL DEFAULT '{}', cursor_json TEXT NOT NULL DEFAULT '{}', phase TEXT NOT NULL DEFAULT 'exporting', status TEXT NOT NULL DEFAULT 'exporting', total_items INTEGER NOT NULL DEFAULT 0, completed_items INTEGER NOT NULL DEFAULT 0, failed_items INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL, started_at TEXT, completed_at TEXT, last_heartbeat_at TEXT, error TEXT ) """, """ CREATE TABLE IF NOT EXISTS export_work_items ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES export_runs(run_id) ON DELETE CASCADE, unit_type TEXT NOT NULL, ordinal INTEGER NOT NULL, document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE, payload_json TEXT NOT NULL DEFAULT '{}', artifact_manifest_json TEXT NOT NULL DEFAULT '{}', status TEXT NOT NULL DEFAULT 'pending', last_error TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, UNIQUE(run_id, ordinal) ) """, "CREATE INDEX IF NOT EXISTS idx_documents_file_hash ON documents(file_hash)", "CREATE INDEX IF NOT EXISTS idx_documents_content_hash ON documents(content_hash)", "CREATE INDEX IF NOT EXISTS idx_documents_lifecycle_status ON documents(lifecycle_status)", "CREATE INDEX IF NOT EXISTS idx_document_source_parts_document_id ON document_source_parts(document_id, part_kind, ordinal)", "CREATE INDEX IF NOT EXISTS idx_previews_document_id ON document_previews(document_id, ordinal)", "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON document_chunks(document_id, chunk_index)", "CREATE INDEX IF NOT EXISTS idx_ingest_runs_status ON ingest_runs(status, phase)", "CREATE INDEX IF NOT EXISTS idx_ingest_runs_cancel_requested ON ingest_runs(cancel_requested_at)", "CREATE INDEX IF NOT EXISTS idx_ingest_work_items_run_status ON ingest_work_items(run_id, status)", "CREATE INDEX IF NOT EXISTS idx_ingest_work_items_commit_order ON ingest_work_items(run_id, commit_order, id)", "CREATE INDEX IF NOT EXISTS idx_ingest_work_items_lease_expires ON ingest_work_items(lease_expires_at)", """ CREATE UNIQUE INDEX IF NOT EXISTS idx_ingest_work_items_source_unique ON ingest_work_items(run_id, unit_type, COALESCE(source_key, ''), COALESCE(rel_path, ''), COALESCE(parent_order, -1)) """, "CREATE INDEX IF NOT EXISTS idx_ingest_prepared_items_run_work ON ingest_prepared_items(run_id, work_item_id)", "CREATE INDEX IF NOT EXISTS idx_ingest_phase_cursors_run_phase ON ingest_phase_cursors(run_id, phase, status)", "CREATE INDEX IF NOT EXISTS idx_ingest_worker_events_run_created ON ingest_worker_events(run_id, created_at)", "CREATE INDEX IF NOT EXISTS idx_ingest_artifact_sweeps_run_state ON ingest_artifact_sweeps(run_id, state)", "CREATE INDEX IF NOT EXISTS idx_entity_rebuild_runs_status ON entity_rebuild_runs(status, phase)", "CREATE INDEX IF NOT EXISTS idx_entity_rebuild_items_run_status ON entity_rebuild_items(run_id, status)", "CREATE INDEX IF NOT EXISTS idx_entity_rebuild_items_lease ON entity_rebuild_items(lease_expires_at)", "CREATE INDEX IF NOT EXISTS idx_export_runs_kind_status ON export_runs(export_kind, status, phase)", "CREATE INDEX IF NOT EXISTS idx_export_work_items_run_status ON export_work_items(run_id, status, ordinal)", ] class RetrieverError(RuntimeError): pass class RetrieverStructuredError(RetrieverError): def __init__(self, message: str, payload: dict[str, object]): super().__init__(message) self.payload = payload def utc_now() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def parse_utc_timestamp(value: object) -> datetime | None: if not value: return None if isinstance(value, datetime): return value.astimezone(timezone.utc) if not isinstance(value, str): return None try: parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError: return None return parsed.astimezone(timezone.utc) def format_utc_timestamp(value: datetime) -> str: return value.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def next_monotonic_utc_timestamp(previous_values: list[object]) -> str: candidate = datetime.now(timezone.utc).replace(microsecond=0) parsed_values = [parsed for parsed in (parse_utc_timestamp(value) for value in previous_values) if parsed is not None] if parsed_values: latest = max(parsed_values) if candidate <= latest: candidate = latest + timedelta(seconds=1) return format_utc_timestamp(candidate) def sha256_file(path: Path) -> str | None: if not path.exists(): return None digest = hashlib.sha256() with path.open("rb") as handle: while True: chunk = handle.read(65536) if not chunk: break digest.update(chunk) return digest.hexdigest() def sha256_text(text: str) -> str: return hashlib.sha256(text.encode("utf-8")).hexdigest() def sha256_bytes(data: bytes) -> str: return hashlib.sha256(data).hexdigest() def sha256_json_value(value: object) -> str: return sha256_text(json.dumps(value, ensure_ascii=True, sort_keys=True, separators=(",", ":"))) def run_command(command: list[str]) -> tuple[bool, str]: try: completed = subprocess.run(command, check=True, capture_output=True, text=True) return True, (completed.stdout or completed.stderr).strip() except Exception as exc: # pragma: no cover - shell probe return False, f"{type(exc).__name__}: {exc}" def workspace_paths(root: Path) -> dict[str, Path]: state_dir = root / ".retriever" tmp_dir = state_dir / "tmp" ingest_tmp_dir = tmp_dir / "ingest" locks_dir = state_dir / "locks" return { "root": root, "state_dir": state_dir, "db_path": state_dir / "retriever.db", "session_path": state_dir / "session.json", "saved_scopes_path": state_dir / "saved_scopes.json", "previews_dir": state_dir / "previews", "text_revisions_dir": state_dir / "text-revisions", "jobs_dir": state_dir / "jobs", "logs_dir": state_dir / "logs", "runtime_path": state_dir / "runtime.json", "tmp_dir": tmp_dir, "ingest_tmp_dir": ingest_tmp_dir, "locks_dir": locks_dir, "ingest_lock_path": locks_dir / "ingest.lock", "entity_rebuild_lock_path": locks_dir / "entity-rebuild.lock", } def ensure_layout(paths: dict[str, Path]) -> None: paths["state_dir"].mkdir(parents=True, exist_ok=True) for key in ( "previews_dir", "text_revisions_dir", "jobs_dir", "logs_dir", "tmp_dir", "ingest_tmp_dir", "locks_dir", ): paths[key].mkdir(parents=True, exist_ok=True) def set_active_workspace_root(root: Path | None) -> None: global ACTIVE_WORKSPACE_ROOT if root is None: ACTIVE_WORKSPACE_ROOT = None return ACTIVE_WORKSPACE_ROOT = Path(root).expanduser().resolve() def active_workspace_paths() -> dict[str, Path] | None: if ACTIVE_WORKSPACE_ROOT is None: return None return workspace_paths(ACTIVE_WORKSPACE_ROOT) def _venv_python_rel_path() -> Path: return Path("Scripts/python.exe") if os.name == "nt" else Path("bin/python") def _canonical_plugin_root_from_tool_path(path: Path) -> Path | None: if ( path.is_file() and path.parent.name == "tool-template" and path.parent.parent.name == "skills" ): return path.parent.parent.parent return None def resolve_plugin_root(root: Path | None = None, *, current_file: str | None = None) -> Path | None: env_root = normalize_whitespace(os.environ.get("RETRIEVER_PLUGIN_ROOT") or "") if env_root: try: return Path(env_root).expanduser().resolve() except OSError: return Path(env_root).expanduser() env_tool = normalize_whitespace(os.environ.get("RETRIEVER_CANONICAL_TOOL_PATH") or "") if env_tool: try: candidate_root = _canonical_plugin_root_from_tool_path(Path(env_tool).expanduser().resolve()) if candidate_root is not None: return candidate_root except OSError: pass canonical_tool = locate_canonical_plugin_tool(current_file=current_file) if canonical_tool is not None: candidate_root = _canonical_plugin_root_from_tool_path(canonical_tool) if candidate_root is not None: return candidate_root candidate_file = current_file or __file__ if candidate_file: try: current_path = Path(candidate_file).resolve() except OSError: current_path = None if current_path is not None: candidate_root = _canonical_plugin_root_from_tool_path(current_path) if candidate_root is not None: return candidate_root if root is not None: try: runtime = read_runtime(workspace_paths(Path(root).expanduser().resolve())["runtime_path"]) except Exception: runtime = None if isinstance(runtime, dict): runtime_payload = runtime.get("plugin_runtime") if isinstance(runtime_payload, dict): plugin_root_value = normalize_whitespace(str(runtime_payload.get("plugin_root") or "")) if plugin_root_value: try: return Path(plugin_root_value).expanduser().resolve() except OSError: return Path(plugin_root_value).expanduser() return None def plugin_runtime_environment_key() -> str: system = re.sub(r"[^a-z0-9]+", "-", platform.system().lower()).strip("-") or "unknown-system" machine = re.sub(r"[^a-z0-9]+", "-", platform.machine().lower()).strip("-") or "unknown-machine" return f"{system}-{machine}-py{sys.version_info.major}.{sys.version_info.minor}" def plugin_runtime_paths(root: Path | None = None, *, current_file: str | None = None) -> dict[str, Path] | None: plugin_root = resolve_plugin_root(root, current_file=current_file) if plugin_root is None: return None runtime_root = plugin_root / ".retriever-plugin-runtime" / plugin_runtime_environment_key() locks_dir = runtime_root / "locks" venv_dir = runtime_root / "venv" return { "plugin_root": plugin_root, "runtime_root": runtime_root, "locks_dir": locks_dir, "venv_dir": venv_dir, "venv_python_path": venv_dir / _venv_python_rel_path(), "requirements_marker_path": runtime_root / ".requirements-version", "install_lock_path": locks_dir / "runtime-install.lock", } def plugin_runtime_site_packages_candidates(paths: dict[str, Path]) -> list[Path]: if os.name == "nt": return [paths["venv_dir"] / "Lib" / "site-packages"] lib_dir = paths["venv_dir"] / "lib" default_path = lib_dir / f"python{sys.version_info.major}.{sys.version_info.minor}" / "site-packages" candidates = [default_path] if lib_dir.exists(): for candidate in sorted(lib_dir.glob("python*/site-packages")): if candidate not in candidates: candidates.append(candidate) return candidates def first_existing_plugin_runtime_site_packages(paths: dict[str, Path]) -> Path | None: for candidate in plugin_runtime_site_packages_candidates(paths): if candidate.exists(): return candidate return None def remove_directory_tree(path: Path) -> bool: try: if not path.exists(): return False shutil.rmtree(path) return True except FileNotFoundError: return False def new_ingest_session_id(now: datetime | None = None) -> str: timestamp = (now or datetime.now(timezone.utc)).strftime("%Y%m%dT%H%M%SZ") return f"{timestamp}-{secrets.token_hex(4)}" def sweep_stale_ingest_tmp_dirs(paths: dict[str, Path]) -> dict[str, object]: ingest_tmp_dir = paths["ingest_tmp_dir"] if not ingest_tmp_dir.exists(): return {"removed": 0, "failures": []} removed = 0 failures: list[dict[str, str]] = [] try: children = sorted(ingest_tmp_dir.iterdir(), key=lambda path: path.name) except OSError as exc: return { "removed": 0, "failures": [{"path": str(ingest_tmp_dir), "error": f"{type(exc).__name__}: {exc}"}], } for child in children: if not child.is_dir(): continue try: if remove_directory_tree(child): removed += 1 except OSError as exc: failures.append( { "path": str(child), "error": f"{type(exc).__name__}: {exc}", } ) return {"removed": removed, "failures": failures} def acquire_os_file_lock(handle) -> None: if os.name == "nt": if msvcrt is None: raise RetrieverError("Windows file locking support is unavailable in this runtime.") handle.seek(0, os.SEEK_END) if handle.tell() == 0: handle.write(b"0") handle.flush() handle.seek(0) msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) return if fcntl is None: raise RetrieverError("POSIX file locking support is unavailable in this runtime.") fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) def release_os_file_lock(handle) -> None: if os.name == "nt": if msvcrt is None: return handle.seek(0) msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) return if fcntl is None: return fcntl.flock(handle.fileno(), fcntl.LOCK_UN) def acquire_workspace_ingest_lock(paths: dict[str, Path]): lock_path = paths["ingest_lock_path"] handle = lock_path.open("a+b") try: acquire_os_file_lock(handle) except RetrieverError: handle.close() raise except OSError as exc: handle.close() if exc.errno in {errno.EACCES, errno.EAGAIN} or isinstance(exc, PermissionError): raise RetrieverError("Another ingest is already running in this workspace. Wait for it to finish and retry.") from exc raise RetrieverError( f"Unable to acquire workspace ingest lock at {lock_path}: {type(exc).__name__}: {exc}" ) from exc return handle def release_workspace_ingest_lock(handle) -> None: try: release_os_file_lock(handle) finally: handle.close() def acquire_workspace_entity_rebuild_lock(paths: dict[str, Path]): lock_path = paths["entity_rebuild_lock_path"] handle = lock_path.open("a+b") try: acquire_os_file_lock(handle) except RetrieverError: handle.close() raise except OSError as exc: handle.close() if exc.errno in {errno.EACCES, errno.EAGAIN} or isinstance(exc, PermissionError): raise RetrieverError( "Another entity rebuild is already running in this workspace. Wait for it to finish and retry." ) from exc raise RetrieverError( f"Unable to acquire workspace entity rebuild lock at {lock_path}: {type(exc).__name__}: {exc}" ) from exc return handle def release_workspace_entity_rebuild_lock(handle) -> None: try: release_os_file_lock(handle) finally: handle.close() @contextmanager def workspace_entity_rebuild_session(paths: dict[str, Path], *, command_name: str): ensure_layout(paths) ingest_lock_handle = acquire_workspace_ingest_lock(paths) entity_lock_handle = None benchmark_mark("workspace_ingest_lock_acquired", command=command_name) try: entity_lock_handle = acquire_workspace_entity_rebuild_lock(paths) benchmark_mark("workspace_entity_rebuild_lock_acquired", command=command_name) yield {"id": new_ingest_session_id()} finally: if entity_lock_handle is not None: release_workspace_entity_rebuild_lock(entity_lock_handle) release_workspace_ingest_lock(ingest_lock_handle) def acquire_plugin_runtime_install_lock(paths: dict[str, Path]): paths["locks_dir"].mkdir(parents=True, exist_ok=True) lock_path = paths["install_lock_path"] handle = lock_path.open("a+b") try: acquire_os_file_lock(handle) except RetrieverError: handle.close() raise except OSError as exc: handle.close() if exc.errno in {errno.EACCES, errno.EAGAIN} or isinstance(exc, PermissionError): raise RetrieverError( "Another shared plugin runtime install is already running. Wait for it to finish and retry." ) from exc raise RetrieverError( f"Unable to acquire plugin runtime install lock at {lock_path}: {type(exc).__name__}: {exc}" ) from exc return handle def release_plugin_runtime_install_lock(handle) -> None: try: release_os_file_lock(handle) finally: handle.close() @contextmanager def workspace_ingest_session(paths: dict[str, Path], *, command_name: str): ensure_layout(paths) lock_handle = acquire_workspace_ingest_lock(paths) benchmark_mark("workspace_ingest_lock_acquired", command=command_name) session_id = new_ingest_session_id() session_dir = paths["ingest_tmp_dir"] / session_id try: stale_tmp_sweep = sweep_stale_ingest_tmp_dirs(paths) stale_tmp_dirs_removed = int(stale_tmp_sweep["removed"]) stale_tmp_dir_failures = list(stale_tmp_sweep.get("failures") or []) warnings = [ f"Could not remove stale ingest tmp dir {failure['path']}: {failure['error']}" for failure in stale_tmp_dir_failures ] session_dir.mkdir(parents=True, exist_ok=True) benchmark_mark( "workspace_ingest_session_ready", command=command_name, session_id=session_id, stale_tmp_dirs_removed=stale_tmp_dirs_removed, stale_tmp_dirs_failed=len(stale_tmp_dir_failures), ) yield { "id": session_id, "tmp_dir": session_dir, "stale_tmp_dirs_removed": stale_tmp_dirs_removed, "stale_tmp_dirs_failed": len(stale_tmp_dir_failures), "warnings": warnings, } finally: try: remove_directory_tree(session_dir) except OSError as exc: benchmark_mark( "workspace_ingest_session_cleanup_failed", command=command_name, session_id=session_id, error=f"{type(exc).__name__}: {exc}", ) release_workspace_ingest_lock(lock_handle) def sqlite_artifact_paths(db_path: Path) -> list[Path]: return [ db_path, Path(f"{db_path}-journal"), Path(f"{db_path}-wal"), Path(f"{db_path}-shm"), ] def stale_sqlite_artifact_paths(db_path: Path) -> list[Path]: db_exists = db_path.exists() sidecars = [path for path in sqlite_artifact_paths(db_path)[1:] if path.exists()] if db_exists: try: if db_path.stat().st_size == 0: return [db_path, *sidecars] except OSError: return [db_path, *sidecars] return [] return sidecars def remove_stale_sqlite_artifacts(db_path: Path) -> list[str]: removed: list[str] = [] for path in stale_sqlite_artifact_paths(db_path): try: path.unlink() removed.append(str(path)) except FileNotFoundError: continue return removed def best_effort_reset_sqlite_artifacts(db_path: Path) -> list[str]: reset_paths: list[str] = [] for path in sqlite_artifact_paths(db_path): if not path.exists(): continue try: if path == db_path: with path.open("wb"): pass reset_paths.append(str(path)) continue path.unlink() reset_paths.append(str(path)) except FileNotFoundError: continue except OSError: if path == db_path: raise try: with path.open("wb"): pass reset_paths.append(str(path)) except OSError: continue return reset_paths def current_journal_mode(connection: sqlite3.Connection) -> str | None: row = connection.execute("PRAGMA journal_mode").fetchone() if row is None or row[0] in (None, ""): return None return str(row[0]).lower() def set_journal_mode(connection: sqlite3.Connection, journal_mode: str) -> str | None: row = connection.execute(f"PRAGMA journal_mode = {journal_mode}").fetchone() if row is None or row[0] in (None, ""): return None return str(row[0]).lower() def configure_supported_sqlite_journal_mode(connection: sqlite3.Connection) -> tuple[str | None, list[str]]: failures: list[str] = [] for requested_mode in ("WAL", "DELETE", "TRUNCATE"): try: journal_mode = set_journal_mode(connection, requested_mode) except sqlite3.DatabaseError as exc: failures.append(f"{requested_mode} failed with {type(exc).__name__}: {exc}") continue normalized_requested_mode = requested_mode.lower() if journal_mode == normalized_requested_mode: return journal_mode, failures failures.append(f"{requested_mode} returned {journal_mode or '<empty>'}") return None, failures def connect_db(db_path: Path) -> sqlite3.Connection: connection = sqlite3.connect(db_path) connection.row_factory = sqlite3.Row connection.execute("PRAGMA foreign_keys = ON") connection.execute("PRAGMA busy_timeout = 5000") journal_mode, journal_mode_failures = configure_supported_sqlite_journal_mode(connection) if journal_mode is None: connection.close() raise RetrieverError( f"Unable to configure SQLite journal mode for {db_path}: " f"{'; '.join(journal_mode_failures)}" ) if journal_mode == "wal": connection.execute("PRAGMA synchronous = NORMAL") return connection def sqlite_bootstrap_seed_required(db_path: Path, error: Exception) -> bool: db_size = file_size_bytes(db_path) if db_size not in (None, 0): return False if isinstance(error, RetrieverError): detail = normalize_whitespace(str(error or "")).lower() return ( "unable to configure sqlite journal mode" in detail or "disk i/o error" in detail or "unable to open database file" in detail or "readonly" in detail or "read-only" in detail or "permission denied" in detail ) return isinstance(error, (sqlite3.DatabaseError, OSError, PermissionError)) def seed_sqlite_db_from_local_temp(db_path: Path) -> dict[str, object]: db_path.parent.mkdir(parents=True, exist_ok=True) reset_artifacts = best_effort_reset_sqlite_artifacts(db_path) seed_file = tempfile.NamedTemporaryFile( prefix="retriever-seed-", suffix=".db", delete=False, ) seed_path = Path(seed_file.name) seed_file.close() seed_journal_mode: str | None = None try: seed_connection = sqlite3.connect(seed_path) try: seed_journal_mode, seed_journal_mode_failures = configure_supported_sqlite_journal_mode(seed_connection) if seed_journal_mode is None: raise RetrieverError( f"Unable to configure SQLite journal mode for temporary seed DB {seed_path}: " f"{'; '.join(seed_journal_mode_failures)}" ) finally: seed_connection.close() with seed_path.open("rb") as src_handle: db_path.write_bytes(src_handle.read()) return { "journal_mode": seed_journal_mode, "reset_artifacts": reset_artifacts, } finally: for artifact in sqlite_artifact_paths(seed_path): try: artifact.unlink() except FileNotFoundError: continue except OSError: continue def file_size_bytes(path: Path) -> int | None: try: return path.stat().st_size if path.exists() else None except OSError: return None def file_mtime_timestamp(path: Path) -> str | None: try: if not path.exists(): return None stat_result = path.stat() return datetime.fromtimestamp( stat_result.st_mtime_ns / 1_000_000_000, timezone.utc, ).isoformat(timespec="microseconds").replace("+00:00", "Z") except OSError: return None def table_info(connection: sqlite3.Connection, table_name: str) -> list[sqlite3.Row]: return connection.execute(f"PRAGMA table_info({quote_identifier(table_name)})").fetchall() def table_columns(connection: sqlite3.Connection, table_name: str) -> set[str]: return {row["name"] for row in table_info(connection, table_name)} def table_exists(connection: sqlite3.Connection, table_name: str) -> bool: row = connection.execute( """ SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ? """, (table_name,), ).fetchone() return row is not None def rename_table_if_needed(connection: sqlite3.Connection, old_name: str, new_name: str) -> bool: if not table_exists(connection, old_name) or table_exists(connection, new_name): return False connection.execute( f"ALTER TABLE {quote_identifier(old_name)} RENAME TO {quote_identifier(new_name)}" ) return True def rename_column_if_needed( connection: sqlite3.Connection, table_name: str, old_name: str, new_name: str, ) -> bool: columns = table_columns(connection, table_name) if old_name not in columns or new_name in columns: return False connection.execute( f"ALTER TABLE {quote_identifier(table_name)} " f"RENAME COLUMN {quote_identifier(old_name)} TO {quote_identifier(new_name)}" ) return True def backfill_legacy_column( connection: sqlite3.Connection, table_name: str, old_name: str, new_name: str, *, treat_blank_as_missing: bool = False, ) -> bool: columns = table_columns(connection, table_name) if old_name not in columns or new_name not in columns: return False before = connection.total_changes where_clause = f"{quote_identifier(new_name)} IS NULL" if treat_blank_as_missing: where_clause = ( f"{quote_identifier(new_name)} IS NULL " f"OR TRIM({quote_identifier(new_name)}) = ''" ) connection.execute( f""" UPDATE {quote_identifier(table_name)} SET {quote_identifier(new_name)} = {quote_identifier(old_name)} WHERE {quote_identifier(old_name)} IS NOT NULL AND ({where_clause}) """ ) return connection.total_changes != before def document_inventory_counts(connection: sqlite3.Connection) -> dict[str, int]: columns = table_columns(connection, "documents") attachment_children_expr = ( "CASE WHEN parent_document_id IS NOT NULL " "AND COALESCE(child_document_kind, 'attachment') = 'attachment' " "AND lifecycle_status != 'deleted' THEN 1 ELSE 0 END" if "child_document_kind" in columns else "CASE WHEN parent_document_id IS NOT NULL AND lifecycle_status != 'deleted' THEN 1 ELSE 0 END" ) row = connection.execute( f""" SELECT COALESCE(SUM(CASE WHEN parent_document_id IS NULL AND lifecycle_status != 'deleted' THEN 1 ELSE 0 END), 0) AS parent_documents, COALESCE(SUM(CASE WHEN parent_document_id IS NULL AND lifecycle_status = 'missing' THEN 1 ELSE 0 END), 0) AS missing_parent_documents, COALESCE(SUM({attachment_children_expr}), 0) AS attachment_children, COALESCE(SUM(CASE WHEN lifecycle_status != 'deleted' THEN 1 ELSE 0 END), 0) AS documents_total FROM documents """ ).fetchone() return { "parent_documents": int(row["parent_documents"]), "missing_parent_documents": int(row["missing_parent_documents"]), "attachment_children": int(row["attachment_children"]), "documents_total": int(row["documents_total"]), } def backfill_source_kinds(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") if not {"source_kind", "production_id", "parent_document_id"}.issubset(columns): return 0 attachment_clause = ( "parent_document_id IS NOT NULL AND COALESCE(child_document_kind, 'attachment') = 'attachment'" if "child_document_kind" in columns else "parent_document_id IS NOT NULL" ) cursor = connection.execute( f""" UPDATE documents SET source_kind = CASE WHEN production_id IS NOT NULL THEN ? WHEN {attachment_clause} THEN ? ELSE ? END WHERE source_kind IS NULL OR TRIM(source_kind) = '' """, (PRODUCTION_SOURCE_KIND, EMAIL_ATTACHMENT_SOURCE_KIND, FILESYSTEM_SOURCE_KIND), ) return int(cursor.rowcount or 0) def ensure_column(connection: sqlite3.Connection, table_name: str, column_definition: str) -> None: column_name = column_definition.split()[0] if column_name in table_columns(connection, table_name): return connection.execute(f"ALTER TABLE {quote_identifier(table_name)} ADD COLUMN {column_definition}") def ensure_export_work_items_nullable_document_id(connection: sqlite3.Connection) -> bool: if not table_exists(connection, "export_work_items"): return False document_id_info = None for row in table_info(connection, "export_work_items"): if row["name"] == "document_id": document_id_info = row break if document_id_info is None or int(document_id_info["notnull"] or 0) == 0: return False backup_table = "export_work_items_document_id_notnull_backup" connection.execute(f"DROP TABLE IF EXISTS {quote_identifier(backup_table)}") connection.execute(f"ALTER TABLE export_work_items RENAME TO {quote_identifier(backup_table)}") connection.execute( """ CREATE TABLE export_work_items ( id INTEGER PRIMARY KEY, run_id TEXT NOT NULL REFERENCES export_runs(run_id) ON DELETE CASCADE, unit_type TEXT NOT NULL, ordinal INTEGER NOT NULL, document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE, payload_json TEXT NOT NULL DEFAULT '{}', artifact_manifest_json TEXT NOT NULL DEFAULT '{}', status TEXT NOT NULL DEFAULT 'pending', last_error TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, UNIQUE(run_id, ordinal) ) """ ) connection.execute( f""" INSERT INTO export_work_items ( id, run_id, unit_type, ordinal, document_id, payload_json, artifact_manifest_json, status, last_error, created_at, updated_at ) SELECT id, run_id, unit_type, ordinal, document_id, payload_json, artifact_manifest_json, status, last_error, created_at, updated_at FROM {quote_identifier(backup_table)} """ ) connection.execute(f"DROP TABLE {quote_identifier(backup_table)}") connection.execute( "CREATE INDEX IF NOT EXISTS idx_export_work_items_run_status ON export_work_items(run_id, status, ordinal)" ) return True def normalize_string_list(raw_value: object) -> list[str]: if raw_value in (None, ""): return [] if isinstance(raw_value, (list, tuple)): values = raw_value else: try: values = json.loads(raw_value) except (TypeError, ValueError, json.JSONDecodeError): return [] if not isinstance(values, list): return [] normalized: list[str] = [] for item in values: if not isinstance(item, str): continue value = item.strip() if value and value not in normalized: normalized.append(value) return normalized def normalize_child_document_kind(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")).lower().replace("-", "_") if not normalized: return None if normalized not in ALLOWED_CHILD_DOCUMENT_KINDS: allowed = ", ".join(sorted(ALLOWED_CHILD_DOCUMENT_KINDS)) raise RetrieverError(f"Unsupported child_document_kind: {value!r}. Expected one of: {allowed}.") return normalized def normalize_conversation_assignment_mode(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")).lower().replace("-", "_") if not normalized: return None if normalized in {CONVERSATION_ASSIGNMENT_MODE_AUTO, CONVERSATION_ASSIGNMENT_MODE_MANUAL}: return normalized raise RetrieverError( "Unsupported conversation_assignment_mode: " f"{value!r}. Expected one of: {CONVERSATION_ASSIGNMENT_MODE_AUTO}, {CONVERSATION_ASSIGNMENT_MODE_MANUAL}." ) def effective_child_document_kind( *, parent_document_id: int | None, child_document_kind: object, ) -> str | None: normalized = normalize_child_document_kind(child_document_kind) if parent_document_id is None: return None return normalized or CHILD_DOCUMENT_KIND_ATTACHMENT def effective_conversation_assignment_mode(conversation_assignment_mode: object) -> str: normalized = normalize_conversation_assignment_mode(conversation_assignment_mode) return normalized or CONVERSATION_ASSIGNMENT_MODE_AUTO def attachment_child_filter_sql(alias: str | None = None) -> str: prefix = f"{alias}." if alias else "" return ( f"{prefix}parent_document_id IS NOT NULL " f"AND COALESCE({prefix}child_document_kind, '{CHILD_DOCUMENT_KIND_ATTACHMENT}') = '{CHILD_DOCUMENT_KIND_ATTACHMENT}'" ) def row_child_document_kind(row: sqlite3.Row | dict[str, object]) -> str | None: if isinstance(row, sqlite3.Row): parent_document_id = row["parent_document_id"] raw_kind = row["child_document_kind"] if "child_document_kind" in row.keys() else None else: parent_document_id = row.get("parent_document_id") raw_kind = row.get("child_document_kind") normalized = normalize_whitespace(str(raw_kind or "")).lower().replace("-", "_") if normalized: return normalized if parent_document_id is not None: return CHILD_DOCUMENT_KIND_ATTACHMENT return None def is_attachment_row(row: sqlite3.Row | dict[str, object]) -> bool: if isinstance(row, sqlite3.Row): parent_document_id = row["parent_document_id"] else: parent_document_id = row.get("parent_document_id") return parent_document_id is not None and row_child_document_kind(row) == CHILD_DOCUMENT_KIND_ATTACHMENT def quote_identifier(name: str) -> str: return '"' + name.replace('"', '""') + '"' def normalize_whitespace(text: str) -> str: text = text.replace("\r\n", "\n").replace("\r", "\n").replace("\x00", "") text = re.sub(r"[ \t]+\n", "\n", text) text = re.sub(r"\n{3,}", "\n\n", text) return text.strip() def normalize_inline_whitespace(text: str) -> str: return re.sub(r"\s+", " ", unicodedata.normalize("NFC", str(text or "")).strip()) def normalize_dataset_name_for_compare(dataset_name: str) -> str: normalized = normalize_inline_whitespace(dataset_name) return normalized.casefold() def normalize_saved_scope_name(scope_name: str) -> str: return normalize_dataset_name_for_compare(scope_name) def normalize_browse_mode(raw_value: object | None) -> str: normalized = normalize_inline_whitespace(str(raw_value or DEFAULT_BROWSE_MODE)).lower() if normalized not in {BROWSE_MODE_DOCUMENTS, BROWSE_MODE_CONVERSATIONS, BROWSE_MODE_ENTITIES}: return DEFAULT_BROWSE_MODE return normalized def default_session_state() -> dict[str, object]: return { "schema_version": SESSION_SCHEMA_VERSION, "scope": {}, "browse_mode": DEFAULT_BROWSE_MODE, "browsing": { BROWSE_MODE_DOCUMENTS: {}, BROWSE_MODE_CONVERSATIONS: {}, BROWSE_MODE_ENTITIES: {}, }, "display": { BROWSE_MODE_DOCUMENTS: {}, BROWSE_MODE_CONVERSATIONS: {}, BROWSE_MODE_ENTITIES: {}, }, } def default_saved_scopes_state() -> dict[str, object]: return { "schema_version": SESSION_SCHEMA_VERSION, "scopes": {}, } def coerce_scope_dataset_entries(raw_value: object) -> list[dict[str, object]]: if not isinstance(raw_value, list): return [] normalized_entries: list[dict[str, object]] = [] seen_ids: set[int] = set() for item in raw_value: if not isinstance(item, dict): continue try: dataset_id = int(item.get("id")) except (TypeError, ValueError): continue dataset_name = normalize_inline_whitespace(str(item.get("name") or "")) if not dataset_name or dataset_id in seen_ids: continue seen_ids.add(dataset_id) normalized_entries.append({"id": dataset_id, "name": dataset_name}) return normalized_entries def coerce_scope_payload(raw_scope: object) -> dict[str, object]: if not isinstance(raw_scope, dict): return {} scope: dict[str, object] = {} keyword = raw_scope.get("keyword") if isinstance(keyword, str) and keyword.strip(): scope["keyword"] = keyword bates = raw_scope.get("bates") if isinstance(bates, dict): begin = normalize_inline_whitespace(str(bates.get("begin") or "")) end = normalize_inline_whitespace(str(bates.get("end") or "")) if begin and end: scope["bates"] = {"begin": begin, "end": end} filter_expression = raw_scope.get("filter") if isinstance(filter_expression, str) and filter_expression.strip(): scope["filter"] = filter_expression dataset_entries = coerce_scope_dataset_entries(raw_scope.get("dataset")) if dataset_entries: scope["dataset"] = dataset_entries from_run_id = raw_scope.get("from_run_id") if from_run_id is not None: try: scope["from_run_id"] = int(from_run_id) except (TypeError, ValueError): pass set_at = raw_scope.get("set_at") if isinstance(set_at, str) and set_at.strip(): scope["set_at"] = set_at return scope def coerce_saved_scope_payload(raw_scope: object) -> dict[str, object]: scope = coerce_scope_payload(raw_scope) saved_at = raw_scope.get("saved_at") if isinstance(raw_scope, dict) else None if isinstance(saved_at, str) and saved_at.strip(): scope["saved_at"] = saved_at scope.pop("set_at", None) return scope def coerce_browsing_sort_payload(raw_value: object) -> list[list[str]]: if not isinstance(raw_value, list): return [] normalized_specs: list[list[str]] = [] for item in raw_value: if not isinstance(item, (list, tuple)) or len(item) != 2: continue field_name = normalize_inline_whitespace(str(item[0] or "")) direction = normalize_inline_whitespace(str(item[1] or "")).lower() if not field_name or direction not in {"asc", "desc"}: continue normalized_specs.append([field_name, direction]) return normalized_specs def coerce_browsing_payload(raw_browsing: object) -> dict[str, object]: if not isinstance(raw_browsing, dict): return {} browsing: dict[str, object] = {} sort_specs = coerce_browsing_sort_payload(raw_browsing.get("sort")) if sort_specs: browsing["sort"] = sort_specs offset = raw_browsing.get("offset") if isinstance(offset, int) and offset >= 0: browsing["offset"] = offset total_known = raw_browsing.get("total_known") if isinstance(total_known, int) and total_known >= 0: browsing["total_known"] = total_known query = raw_browsing.get("query") if isinstance(query, str) and query.strip(): browsing["query"] = normalize_whitespace(query) include_ignored = raw_browsing.get("include_ignored") if isinstance(include_ignored, bool): browsing["include_ignored"] = include_ignored run_at = raw_browsing.get("run_at") if isinstance(run_at, str) and run_at.strip(): browsing["run_at"] = run_at return browsing def coerce_display_payload(raw_display: object) -> dict[str, object]: if not isinstance(raw_display, dict): return {} display: dict[str, object] = {} columns = raw_display.get("columns") if isinstance(columns, list): normalized_columns = [normalize_inline_whitespace(str(value)) for value in columns if normalize_inline_whitespace(str(value))] if normalized_columns: display["columns"] = normalized_columns page_size = raw_display.get("page_size") if isinstance(page_size, int) and page_size > 0: display["page_size"] = page_size return display def coerce_mode_payloads(raw_value: object, payload_coercer) -> dict[str, object]: normalized_payloads = { BROWSE_MODE_DOCUMENTS: {}, BROWSE_MODE_CONVERSATIONS: {}, BROWSE_MODE_ENTITIES: {}, } if not isinstance(raw_value, dict): return normalized_payloads if any( key in raw_value for key in ("columns", "page_size", "sort", "offset", "total_known", "query", "include_ignored", "run_at") ): normalized_payloads[BROWSE_MODE_DOCUMENTS] = payload_coercer(raw_value) return normalized_payloads for browse_mode in (BROWSE_MODE_DOCUMENTS, BROWSE_MODE_CONVERSATIONS, BROWSE_MODE_ENTITIES): normalized_payloads[browse_mode] = payload_coercer(raw_value.get(browse_mode)) return normalized_payloads def coerce_session_state(raw_value: object) -> dict[str, object]: session = default_session_state() if not isinstance(raw_value, dict): return session schema_version = raw_value.get("schema_version") if isinstance(schema_version, int): session["schema_version"] = schema_version session["scope"] = coerce_scope_payload(raw_value.get("scope")) session["browse_mode"] = normalize_browse_mode(raw_value.get("browse_mode")) session["browsing"] = coerce_mode_payloads(raw_value.get("browsing"), coerce_browsing_payload) session["display"] = coerce_mode_payloads(raw_value.get("display"), coerce_display_payload) return session def coerce_saved_scopes_state(raw_value: object) -> dict[str, object]: saved_scopes = default_saved_scopes_state() if not isinstance(raw_value, dict): return saved_scopes schema_version = raw_value.get("schema_version") if isinstance(schema_version, int): saved_scopes["schema_version"] = schema_version scopes = raw_value.get("scopes") if not isinstance(scopes, dict): return saved_scopes normalized_scopes: dict[str, object] = {} for scope_name, scope_payload in scopes.items(): if not isinstance(scope_name, str): continue normalized_name = normalize_saved_scope_name(scope_name) if not normalized_name: continue normalized_scopes[scope_name] = coerce_saved_scope_payload(scope_payload) saved_scopes["scopes"] = normalized_scopes return saved_scopes def read_json_state(path: Path, default_factory) -> dict[str, object]: if not path.exists(): return default_factory() try: payload = json.loads(path.read_text(encoding="utf-8")) except (OSError, ValueError, json.JSONDecodeError) as exc: raise RetrieverError(f"Could not read state file {path}: {exc}") from exc coerced = default_factory() if default_factory is default_session_state: coerced = coerce_session_state(payload) elif default_factory is default_saved_scopes_state: coerced = coerce_saved_scopes_state(payload) return coerced def write_json_atomic(path: Path, payload: dict[str, object]) -> None: path.parent.mkdir(parents=True, exist_ok=True) file_handle = tempfile.NamedTemporaryFile( mode="w", encoding="utf-8", dir=str(path.parent), prefix=f".{path.name}.", suffix=".tmp", delete=False, ) temp_path = Path(file_handle.name) try: with file_handle: json.dump(payload, file_handle, indent=2, sort_keys=True) file_handle.write("\n") file_handle.flush() temp_path.replace(path) except Exception: try: temp_path.unlink() except FileNotFoundError: pass raise def read_session_state(paths: dict[str, Path]) -> dict[str, object]: return read_json_state(paths["session_path"], default_session_state) def write_session_state(paths: dict[str, Path], payload: dict[str, object]) -> None: coerced = coerce_session_state(payload) coerced["schema_version"] = SESSION_SCHEMA_VERSION write_json_atomic(paths["session_path"], coerced) def read_saved_scopes_state(paths: dict[str, Path]) -> dict[str, object]: return read_json_state(paths["saved_scopes_path"], default_saved_scopes_state) def write_saved_scopes_state(paths: dict[str, Path], payload: dict[str, object]) -> None: coerced = coerce_saved_scopes_state(payload) coerced["schema_version"] = SESSION_SCHEMA_VERSION write_json_atomic(paths["saved_scopes_path"], coerced) def normalize_extension(path: Path) -> str: return path.suffix.lower().lstrip(".") def format_control_number( batch_number: int, family_sequence: int, attachment_sequence: int | None = None, ) -> str: base = f"{CONTROL_NUMBER_PREFIX}{batch_number:0{CONTROL_NUMBER_BATCH_WIDTH}d}.{family_sequence:0{CONTROL_NUMBER_FAMILY_WIDTH}d}" if attachment_sequence is None: return base return f"{base}.{attachment_sequence:0{CONTROL_NUMBER_ATTACHMENT_WIDTH}d}" def parse_control_number(control_number: object) -> tuple[int, int, int | None] | None: if not isinstance(control_number, str): return None match = re.fullmatch( rf"{CONTROL_NUMBER_PREFIX}(\d{{{CONTROL_NUMBER_BATCH_WIDTH}}})\.(\d{{{CONTROL_NUMBER_FAMILY_WIDTH}}})(?:\.(\d{{{CONTROL_NUMBER_ATTACHMENT_WIDTH}}}))?", control_number.strip(), ) if not match: return None attachment_sequence = int(match.group(3)) if match.group(3) is not None else None return int(match.group(1)), int(match.group(2)), attachment_sequence def parse_bates_identifier(value: object) -> dict[str, object] | None: if not isinstance(value, str): return None normalized = value.strip() if not normalized: return None match = re.fullmatch(r"(?P<prefix>.*?)(?P<number>\d+)$", normalized) if match is None: return None prefix = match.group("prefix") number_text = match.group("number") return { "raw": normalized, "prefix": prefix, "prefix_normalized": prefix.strip().upper(), "number": int(number_text), "width": len(number_text), } def bates_series_key(parsed: dict[str, object] | None) -> tuple[str, int] | None: if not parsed: return None return (str(parsed["prefix_normalized"]), int(parsed["width"])) def bates_range_compatible(left: dict[str, object] | None, right: dict[str, object] | None) -> bool: left_key = bates_series_key(left) right_key = bates_series_key(right) return left_key is not None and left_key == right_key def bates_inclusive_contains( begin_value: object, end_value: object, query_value: object, ) -> bool: begin = parse_bates_identifier(begin_value) end = parse_bates_identifier(end_value) query = parse_bates_identifier(query_value) if not bates_range_compatible(begin, end) or not bates_range_compatible(begin, query): return False assert begin is not None and end is not None and query is not None return int(begin["number"]) <= int(query["number"]) <= int(end["number"]) def bates_ranges_overlap( begin_value: object, end_value: object, query_begin_value: object, query_end_value: object, ) -> bool: begin = parse_bates_identifier(begin_value) end = parse_bates_identifier(end_value) query_begin = parse_bates_identifier(query_begin_value) query_end = parse_bates_identifier(query_end_value) if not all((bates_range_compatible(begin, end), bates_range_compatible(query_begin, query_end), bates_range_compatible(begin, query_begin))): return False assert begin is not None and end is not None and query_begin is not None and query_end is not None return int(begin["number"]) <= int(query_end["number"]) and int(end["number"]) >= int(query_begin["number"]) def bates_sort_key(value: object) -> tuple[int, str, int, str]: parsed = parse_bates_identifier(value) if parsed is None: return (1, "", 0, str(value or "")) return (0, str(parsed["prefix_normalized"]), int(parsed["number"]), str(parsed["raw"])) def parse_bates_query(query: str) -> tuple[str, str] | tuple[None, None]: stripped = query.strip() if not stripped: return None, None range_match = re.fullmatch(r"\s*(\S+)\s*[-–]\s*(\S+)\s*", stripped) if range_match: left = range_match.group(1) right = range_match.group(2) if parse_bates_identifier(left) and parse_bates_identifier(right): return left, right if " " not in stripped and parse_bates_identifier(stripped): return stripped, stripped return None, None def normalize_internal_rel_path(path: Path) -> str: normalized = path.as_posix() while normalized.startswith("./"): normalized = normalized[2:] return normalized.lstrip("/") INTERNAL_REL_PATH_PREFIX = "_retriever" def is_internal_rel_path(rel_path: str | None) -> bool: if not rel_path: return False return normalize_internal_rel_path(Path(rel_path)).startswith(f"{INTERNAL_REL_PATH_PREFIX}/") def document_absolute_path(paths: dict[str, Path], rel_path: str | None) -> Path: """Resolve a documents.rel_path (possibly internal) to its absolute location. Internal rel_paths (those beginning with ``_retriever/``) address files that live under the workspace's ``.retriever`` state directory. Regular rel_paths are relative to the workspace root. """ text = str(rel_path or "").strip() if not text: return paths["root"] path = Path(text) if path.parts and path.parts[0] == INTERNAL_REL_PATH_PREFIX: state_relative = Path(*path.parts[1:]) if len(path.parts) > 1 else Path() return paths["state_dir"] / state_relative return paths["root"] / text def normalize_source_item_id(value: object) -> str: text = normalize_whitespace(str(value or "")) if not text: raise RetrieverError("Container-derived documents require a stable source item id.") return text def encode_source_item_id_for_path(source_item_id: str) -> str: encoded = base64.urlsafe_b64encode(normalize_source_item_id(source_item_id).encode("utf-8")).decode("ascii") return encoded.rstrip("=") or "item" def container_source_rel_path_from_message_rel_path(rel_path: str | None) -> str | None: if not rel_path: return None normalized = normalize_internal_rel_path(Path(rel_path)) parts = Path(normalized).parts if len(parts) < 5 or parts[0] != INTERNAL_REL_PATH_PREFIX or parts[1] != "sources": return None try: messages_index = parts.index("messages") except ValueError: return None if messages_index <= 2: return None return Path(*parts[2:messages_index]).as_posix() def infer_source_custodian( *, source_kind: str | None, source_rel_path: str | None, parent_custodian: str | None = None, ) -> str | None: inherited = normalize_whitespace(str(parent_custodian or "")) if inherited: return inherited normalized_source_kind = normalize_whitespace(str(source_kind or "")).lower() normalized_source_rel_path = normalize_whitespace(str(source_rel_path or "")) if normalized_source_kind == MBOX_SOURCE_KIND and normalized_source_rel_path: return mbox_custodian_email_from_source_rel_path(normalized_source_rel_path) if normalized_source_kind == PST_SOURCE_KIND and normalized_source_rel_path: basename = normalize_whitespace(Path(normalized_source_rel_path).stem) if not basename: return None basename_email = normalize_entity_email(basename) if basename_email: entity_type = entity_type_from_candidate_parts(name_value=None, email_value=basename_email) if entity_type in {ENTITY_TYPE_SHARED_MAILBOX, ENTITY_TYPE_SYSTEM_MAILBOX}: return None return basename_email stripped = normalize_whitespace(strip_entity_container_words(basename)) return stripped or None return None GOOGLE_VAULT_MBOX_BASENAME_PATTERN = re.compile( r"^(?P<label>.+?)--(?P<email>[^-]+@[^-]+\.[^-]+)-(?P<random>[A-Za-z0-9_-]{4,})$" ) GOOGLE_TAKEOUT_MBOX_BASENAMES = { "all mail", "all mail including spam and trash", "archive", "chats", "drafts", "important", "inbox", "sent", "spam", "starred", "trash", } def parse_google_vault_mbox_basename(basename: object) -> dict[str, str] | None: normalized = normalize_whitespace(str(basename or "")) if not normalized: return None match = GOOGLE_VAULT_MBOX_BASENAME_PATTERN.match(normalized) if match is None: return None email = normalize_entity_email(match.group("email")) if not email: return None return { "label": normalize_whitespace(match.group("label")), "email": email, "random": match.group("random"), } def is_google_takeout_mbox_basename(basename: object) -> bool: normalized = normalize_entity_lookup_text(basename) return normalized in GOOGLE_TAKEOUT_MBOX_BASENAMES def mbox_custodian_email_from_source_rel_path(source_rel_path: object) -> str | None: basename = normalize_whitespace(Path(normalize_whitespace(str(source_rel_path or ""))).stem) if not basename: return None vault_parts = parse_google_vault_mbox_basename(basename) if vault_parts is not None: return vault_parts["email"] if is_google_takeout_mbox_basename(basename): return None return None def filesystem_dataset_locator() -> str: return "." def filesystem_dataset_name(root: Path | None = None) -> str: if root is not None: candidate = normalize_whitespace(root.resolve().name) if candidate: return candidate return "Workspace files" def container_dataset_name(source_rel_path: str, fallback_label: str) -> str: candidate = normalize_whitespace(Path(source_rel_path).name) return candidate or normalize_whitespace(source_rel_path) or fallback_label def pst_dataset_name(source_rel_path: str) -> str: return container_dataset_name(source_rel_path, "PST Dataset") def mbox_dataset_name(source_rel_path: str) -> str: return container_dataset_name(source_rel_path, "MBOX Dataset") def slack_export_dataset_name(source_rel_path: str) -> str: candidate = normalize_whitespace(source_rel_path) if candidate: return f"Slack Export: {candidate}" return "Slack Export" def production_dataset_name(rel_root: str, production_name: str | None = None) -> str: preferred = normalize_whitespace(str(production_name or "")) if preferred: return preferred candidate = normalize_whitespace(Path(rel_root).name) return candidate or normalize_whitespace(rel_root) or "Production Dataset" def dataset_source_absolute_path(root: Path, source_locator: str | None) -> Path | None: normalized_source_locator = normalize_whitespace(str(source_locator or "")) if not normalized_source_locator or normalized_source_locator == filesystem_dataset_locator(): return None return root / normalized_source_locator def manual_dataset_locator(dataset_name: str | None = None) -> str: seed = normalize_whitespace(str(dataset_name or "")) or "dataset" return f"manual:{sha256_text(f'{seed}:{utc_now()}')[:16]}" def normalized_dataset_name_or_default(dataset_name: str | None, default: str = "Dataset") -> str: normalized = normalize_inline_whitespace(str(dataset_name or "")) return normalized or default def get_dataset_row( connection: sqlite3.Connection, *, source_kind: str, dataset_locator: str, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM datasets WHERE source_kind = ? AND dataset_locator = ? """, (source_kind, dataset_locator), ).fetchone() def ensure_dataset_row( connection: sqlite3.Connection, *, source_kind: str, dataset_locator: str, dataset_name: str, ) -> int: now = utc_now() normalized_dataset_name = normalized_dataset_name_or_default(dataset_name) normalized_compare_name = normalize_dataset_name_for_compare(normalized_dataset_name) existing_row = get_dataset_row( connection, source_kind=source_kind, dataset_locator=dataset_locator, ) if existing_row is None: external_id_auto_merge_names = ( ["slack_user_id"] if normalize_whitespace(str(source_kind or "")).lower() == SLACK_EXPORT_SOURCE_KIND else [] ) try: connection.execute( """ INSERT INTO datasets ( source_kind, dataset_locator, dataset_name, dataset_name_normalized, external_id_auto_merge_names_json, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( source_kind, dataset_locator, normalized_dataset_name, normalized_compare_name, json.dumps(external_id_auto_merge_names, ensure_ascii=True, sort_keys=True), now, now, ), ) except sqlite3.IntegrityError as exc: conflicting_row = connection.execute( """ SELECT id, dataset_name FROM datasets WHERE dataset_name_normalized = ? ORDER BY id ASC LIMIT 1 """, (normalized_compare_name,), ).fetchone() if conflicting_row is not None: raise RetrieverError( f"Dataset name {normalized_dataset_name!r} is already used by dataset {conflicting_row['id']} ({conflicting_row['dataset_name']!r})." ) from exc raise return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) existing_name = normalized_dataset_name_or_default(str(existing_row["dataset_name"] or "")) existing_normalized = normalize_inline_whitespace(str(existing_row["dataset_name_normalized"] or "")) if existing_normalized != normalize_dataset_name_for_compare(existing_name): connection.execute( """ UPDATE datasets SET dataset_name = ?, dataset_name_normalized = ?, updated_at = ? WHERE id = ? """, (existing_name, normalize_dataset_name_for_compare(existing_name), now, existing_row["id"]), ) return int(existing_row["id"]) def create_dataset_row( connection: sqlite3.Connection, dataset_name: str, *, source_kind: str | None = None, dataset_locator: str | None = None, ) -> int: normalized_name = normalized_dataset_name_or_default(dataset_name) normalized_source_kind = normalize_whitespace(str(source_kind or MANUAL_DATASET_SOURCE_KIND)).lower() normalized_locator = normalize_whitespace(str(dataset_locator or "")) if not normalized_locator: normalized_locator = manual_dataset_locator(normalized_name) return ensure_dataset_row( connection, source_kind=normalized_source_kind or MANUAL_DATASET_SOURCE_KIND, dataset_locator=normalized_locator, dataset_name=normalized_name, ) def get_dataset_source_row( connection: sqlite3.Connection, *, source_kind: str, source_locator: str, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM dataset_sources WHERE source_kind = ? AND source_locator = ? """, (source_kind, source_locator), ).fetchone() def ensure_dataset_source_row( connection: sqlite3.Connection, *, dataset_id: int, source_kind: str, source_locator: str, ) -> int: normalized_source_kind = normalize_whitespace(source_kind).lower() normalized_source_locator = normalize_whitespace(source_locator) if not normalized_source_kind or not normalized_source_locator: raise RetrieverError("Dataset sources require non-empty source_kind and source_locator.") now = utc_now() existing_row = get_dataset_source_row( connection, source_kind=normalized_source_kind, source_locator=normalized_source_locator, ) if existing_row is None: connection.execute( """ INSERT INTO dataset_sources ( dataset_id, source_kind, source_locator, created_at, updated_at ) VALUES (?, ?, ?, ?, ?) """, (dataset_id, normalized_source_kind, normalized_source_locator, now, now), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) if int(existing_row["dataset_id"]) != dataset_id: raise RetrieverError( f"Source {normalized_source_kind}:{normalized_source_locator} is already bound to dataset {existing_row['dataset_id']}." ) connection.execute( """ UPDATE dataset_sources SET updated_at = ? WHERE id = ? """, (now, existing_row["id"]), ) return int(existing_row["id"]) def ensure_source_backed_dataset( connection: sqlite3.Connection, *, source_kind: str, source_locator: str, dataset_name: str, ) -> tuple[int, int]: normalized_source_kind = normalize_whitespace(source_kind).lower() normalized_source_locator = normalize_whitespace(source_locator) existing_source = get_dataset_source_row( connection, source_kind=normalized_source_kind, source_locator=normalized_source_locator, ) if existing_source is not None: return int(existing_source["dataset_id"]), int(existing_source["id"]) dataset_id = ensure_dataset_row( connection, source_kind=normalized_source_kind, dataset_locator=normalized_source_locator, dataset_name=normalize_whitespace(dataset_name) or "Dataset", ) dataset_source_id = ensure_dataset_source_row( connection, dataset_id=dataset_id, source_kind=normalized_source_kind, source_locator=normalized_source_locator, ) return dataset_id, dataset_source_id def get_conversation_row( connection: sqlite3.Connection, *, source_kind: str, source_locator: str, conversation_key: str, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM conversations WHERE source_kind = ? AND source_locator = ? AND conversation_key = ? """, ( normalize_whitespace(source_kind).lower(), normalize_whitespace(source_locator), normalize_whitespace(conversation_key), ), ).fetchone() def upsert_conversation_row( connection: sqlite3.Connection, *, source_kind: str, source_locator: str, conversation_key: str, conversation_type: str, display_name: str, ) -> int: normalized_source_kind = normalize_whitespace(source_kind).lower() normalized_source_locator = normalize_whitespace(source_locator) normalized_conversation_key = normalize_whitespace(conversation_key) normalized_conversation_type = normalize_whitespace(conversation_type).lower() normalized_display_name = normalize_whitespace(display_name) if not all( ( normalized_source_kind, normalized_source_locator, normalized_conversation_key, normalized_conversation_type, normalized_display_name, ) ): raise RetrieverError("Conversations require non-empty source kind, source locator, key, type, and display name.") existing_row = get_conversation_row( connection, source_kind=normalized_source_kind, source_locator=normalized_source_locator, conversation_key=normalized_conversation_key, ) now = utc_now() if existing_row is None: connection.execute( """ INSERT INTO conversations ( source_kind, source_locator, conversation_key, conversation_type, display_name, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( normalized_source_kind, normalized_source_locator, normalized_conversation_key, normalized_conversation_type, normalized_display_name, now, now, ), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) connection.execute( """ UPDATE conversations SET conversation_type = ?, display_name = ?, updated_at = ? WHERE id = ? """, (normalized_conversation_type, normalized_display_name, now, int(existing_row["id"])), ) return int(existing_row["id"]) def ensure_dataset_document_membership( connection: sqlite3.Connection, *, dataset_id: int, document_id: int, dataset_source_id: int | None = None, ) -> int: now = utc_now() if dataset_source_id is None: existing_row = connection.execute( """ SELECT id FROM dataset_documents WHERE dataset_id = ? AND document_id = ? AND dataset_source_id IS NULL """, (dataset_id, document_id), ).fetchone() else: existing_row = connection.execute( """ SELECT id FROM dataset_documents WHERE dataset_id = ? AND document_id = ? AND dataset_source_id = ? """, (dataset_id, document_id, dataset_source_id), ).fetchone() if existing_row is None: connection.execute( """ INSERT INTO dataset_documents ( dataset_id, document_id, dataset_source_id, created_at, updated_at ) VALUES (?, ?, ?, ?, ?) """, (dataset_id, document_id, dataset_source_id, now, now), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) return int(existing_row["id"]) def get_dataset_row_by_id(connection: sqlite3.Connection, dataset_id: int) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM datasets WHERE id = ? """, (dataset_id,), ).fetchone() def find_dataset_rows_by_name(connection: sqlite3.Connection, dataset_name: str) -> list[sqlite3.Row]: normalized_name = normalize_dataset_name_for_compare(dataset_name) if not normalized_name: return [] return connection.execute( """ SELECT * FROM datasets WHERE dataset_name_normalized = ? ORDER BY id ASC """, (normalized_name,), ).fetchall() def resolve_dataset_row( connection: sqlite3.Connection, *, dataset_id: int | None = None, dataset_name: str | None = None, ) -> sqlite3.Row: if dataset_id is None and dataset_name is None: raise RetrieverError("Specify either dataset_id or dataset_name.") if dataset_id is not None: row = get_dataset_row_by_id(connection, dataset_id) if row is None: raise RetrieverError(f"Unknown dataset id: {dataset_id}") if dataset_name is not None and normalize_dataset_name_for_compare(str(row["dataset_name"] or "")) != normalize_dataset_name_for_compare(dataset_name): raise RetrieverError( f"Dataset id {dataset_id} is named {row['dataset_name']!r}, not {normalize_inline_whitespace(dataset_name)!r}." ) return row matches = find_dataset_rows_by_name(connection, str(dataset_name or "")) if not matches: raise RetrieverError(f"Unknown dataset name: {dataset_name}") return matches[0] def rename_dataset_row( connection: sqlite3.Connection, dataset_id: int, new_dataset_name: str, root: Path | None = None, ) -> dict[str, object]: dataset_row = get_dataset_row_by_id(connection, dataset_id) if dataset_row is None: raise RetrieverError(f"Unknown dataset id: {dataset_id}") normalized_name = normalized_dataset_name_or_default(new_dataset_name) normalized_compare_name = normalize_dataset_name_for_compare(normalized_name) conflicting_row = connection.execute( """ SELECT id, dataset_name FROM datasets WHERE dataset_name_normalized = ? AND id != ? ORDER BY id ASC LIMIT 1 """, (normalized_compare_name, dataset_id), ).fetchone() if conflicting_row is not None: raise RetrieverError( f"Dataset name {normalized_name!r} is already used by dataset {conflicting_row['id']} ({conflicting_row['dataset_name']!r})." ) now = utc_now() connection.execute( """ UPDATE datasets SET dataset_name = ?, dataset_name_normalized = ?, updated_at = ? WHERE id = ? """, (normalized_name, normalized_compare_name, now, dataset_id), ) return dataset_summary_by_id(connection, dataset_id, root=root) def refresh_document_dataset_cache(connection: sqlite3.Connection, document_id: int) -> int | None: membership_rows = connection.execute( """ SELECT DISTINCT dataset_id FROM dataset_documents WHERE document_id = ? ORDER BY dataset_id ASC """, (document_id,), ).fetchall() cached_dataset_id = int(membership_rows[0]["dataset_id"]) if len(membership_rows) == 1 else None connection.execute( """ UPDATE documents SET dataset_id = ? WHERE id = ? """, (cached_dataset_id, document_id), ) return cached_dataset_id def dataset_container_size_bytes( root: Path | None, dataset_row: sqlite3.Row, source_bindings: list[dict[str, object]], ) -> int | None: if root is None: return None normalized_source_kind = normalize_whitespace(str(dataset_row["source_kind"] or "")).lower() if normalized_source_kind not in {MBOX_SOURCE_KIND, PST_SOURCE_KIND}: return None candidate_locators: list[str] = [] dataset_locator = normalize_whitespace(str(dataset_row["dataset_locator"] or "")) if dataset_locator: candidate_locators.append(dataset_locator) for binding in source_bindings: source_locator = normalize_whitespace(str(binding.get("source_locator") or "")) if source_locator and source_locator not in candidate_locators: candidate_locators.append(source_locator) for source_locator in candidate_locators: source_path = dataset_source_absolute_path(root, source_locator) if source_path is None: continue size_bytes = file_size_bytes(source_path) if size_bytes is not None: return size_bytes return None def list_dataset_summaries(connection: sqlite3.Connection, root: Path | None = None) -> list[dict[str, object]]: dataset_rows = connection.execute( """ SELECT * FROM datasets ORDER BY LOWER(dataset_name) ASC, id ASC """ ).fetchall() if not dataset_rows: return [] dataset_ids = [int(row["id"]) for row in dataset_rows] placeholders = ", ".join("?" for _ in dataset_ids) membership_rows = connection.execute( f""" SELECT dataset_id, COUNT(DISTINCT document_id) AS document_count, COUNT(DISTINCT CASE WHEN dataset_source_id IS NULL THEN document_id END) AS manual_document_count, COUNT(DISTINCT CASE WHEN dataset_source_id IS NOT NULL THEN document_id END) AS source_document_count FROM dataset_documents WHERE dataset_id IN ({placeholders}) GROUP BY dataset_id """, dataset_ids, ).fetchall() membership_counts = { int(row["dataset_id"]): { "document_count": int(row["document_count"] or 0), "manual_document_count": int(row["manual_document_count"] or 0), "source_document_count": int(row["source_document_count"] or 0), } for row in membership_rows } document_rows = connection.execute( f""" WITH distinct_dataset_documents AS ( SELECT DISTINCT dataset_id, document_id FROM dataset_documents WHERE dataset_id IN ({placeholders}) ) SELECT distinct_dataset_documents.dataset_id, d.id AS document_id, d.file_size, d.content_type, d.custodians_json, d.date_created, d.date_modified FROM distinct_dataset_documents JOIN documents d ON d.id = distinct_dataset_documents.document_id ORDER BY distinct_dataset_documents.dataset_id ASC, d.id ASC """, dataset_ids, ).fetchall() stats_by_dataset: dict[int, dict[str, object]] = defaultdict( lambda: { "size_bytes": 0, "sized_document_count": 0, "content_type_counts": defaultdict(int), "custodians": set(), "time_range_start": None, "time_range_end": None, } ) for row in document_rows: dataset_id = int(row["dataset_id"]) stats = stats_by_dataset[dataset_id] file_size_value: int | None = None if row["file_size"] is not None: try: file_size_value = int(row["file_size"]) except (TypeError, ValueError): file_size_value = None if file_size_value is not None and file_size_value >= 0: stats["size_bytes"] = int(stats["size_bytes"]) + file_size_value stats["sized_document_count"] = int(stats["sized_document_count"]) + 1 content_type = normalize_whitespace(str(row["content_type"] or "")) or "Unknown" content_type_counts = stats["content_type_counts"] content_type_counts[content_type] += 1 custodians = stats["custodians"] custodians.update(parse_document_custodians_json(row["custodians_json"])) start_candidate = normalize_datetime(row["date_created"]) or normalize_datetime(row["date_modified"]) if start_candidate is not None: current_start = stats["time_range_start"] if current_start is None or start_candidate < current_start: stats["time_range_start"] = start_candidate end_candidate = normalize_datetime(row["date_modified"]) or normalize_datetime(row["date_created"]) if end_candidate is not None: current_end = stats["time_range_end"] if current_end is None or end_candidate > current_end: stats["time_range_end"] = end_candidate source_rows = connection.execute( f""" SELECT * FROM dataset_sources WHERE dataset_id IN ({placeholders}) ORDER BY dataset_id ASC, source_kind ASC, source_locator ASC, id ASC """, dataset_ids, ).fetchall() sources_by_dataset: dict[int, list[dict[str, object]]] = defaultdict(list) for row in source_rows: sources_by_dataset[int(row["dataset_id"])].append( { "id": int(row["id"]), "source_kind": row["source_kind"], "source_locator": row["source_locator"], "created_at": row["created_at"], "updated_at": row["updated_at"], } ) summaries: list[dict[str, object]] = [] for row in dataset_rows: dataset_id = int(row["id"]) counts = membership_counts.get( dataset_id, {"document_count": 0, "manual_document_count": 0, "source_document_count": 0}, ) dataset_stats = stats_by_dataset.get(dataset_id) source_bindings = sources_by_dataset.get(dataset_id, []) if dataset_stats is None: size_bytes = None sized_document_count = 0 size_basis = None content_types: list[dict[str, object]] = [] custodians: list[str] = [] time_range_start = None time_range_end = None else: sized_document_count = int(dataset_stats["sized_document_count"] or 0) size_bytes = int(dataset_stats["size_bytes"] or 0) if sized_document_count else None size_basis = "documents" if size_bytes is not None else None content_types = [ {"name": name, "count": count} for name, count in sorted( dataset_stats["content_type_counts"].items(), key=lambda item: (-int(item[1]), str(item[0]).lower(), str(item[0])), ) ] custodians = sorted( [str(value) for value in dataset_stats["custodians"]], key=lambda value: (value.lower(), value), ) time_range_start = dataset_stats["time_range_start"] time_range_end = dataset_stats["time_range_end"] container_size_bytes = dataset_container_size_bytes(root, row, source_bindings) if container_size_bytes is not None: size_bytes = container_size_bytes size_basis = "container" summaries.append( { "id": dataset_id, "dataset_name": row["dataset_name"], "source_kind": row["source_kind"], "dataset_locator": row["dataset_locator"], "created_at": row["created_at"], "updated_at": row["updated_at"], "document_count": counts["document_count"], "manual_document_count": counts["manual_document_count"], "source_document_count": counts["source_document_count"], "size_bytes": size_bytes, "size_basis": size_basis, "sized_document_count": sized_document_count, "custodians": custodians, "content_types": content_types, "time_range_start": time_range_start, "time_range_end": time_range_end, "source_binding_count": len(source_bindings), "source_bindings": source_bindings, "merge_policy": dataset_merge_policy_payload_from_row(row), } ) return summaries def dataset_summary_by_id(connection: sqlite3.Connection, dataset_id: int, root: Path | None = None) -> dict[str, object]: for summary in list_dataset_summaries(connection, root=root): if int(summary["id"]) == dataset_id: return summary raise RetrieverError(f"Unknown dataset id: {dataset_id}") def add_documents_to_dataset( connection: sqlite3.Connection, *, dataset_id: int, document_ids: list[int], ) -> dict[str, list[int]]: if not document_ids: return {"added_document_ids": [], "already_present_document_ids": []} unique_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) placeholders = ", ".join("?" for _ in unique_document_ids) existing_rows = connection.execute( f""" SELECT id FROM documents WHERE id IN ({placeholders}) """, unique_document_ids, ).fetchall() existing_ids = {int(row["id"]) for row in existing_rows} missing_ids = [document_id for document_id in unique_document_ids if document_id not in existing_ids] if missing_ids: raise RetrieverError(f"Unknown document ids: {', '.join(str(document_id) for document_id in missing_ids)}") current_rows = connection.execute( f""" SELECT DISTINCT document_id FROM dataset_documents WHERE dataset_id = ? AND document_id IN ({placeholders}) """, [dataset_id, *unique_document_ids], ).fetchall() current_ids = {int(row["document_id"]) for row in current_rows} added_document_ids: list[int] = [] already_present_document_ids: list[int] = [] for document_id in unique_document_ids: if document_id in current_ids: already_present_document_ids.append(document_id) continue ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=None, ) refresh_document_dataset_cache(connection, document_id) added_document_ids.append(document_id) if added_document_ids: connection.execute( """ UPDATE datasets SET updated_at = ? WHERE id = ? """, (utc_now(), dataset_id), ) return { "added_document_ids": added_document_ids, "already_present_document_ids": already_present_document_ids, } def remove_documents_from_dataset( connection: sqlite3.Connection, *, dataset_id: int, document_ids: list[int], ) -> dict[str, list[int]]: if not document_ids: return { "removed_document_ids": [], "not_present_document_ids": [], "documents_without_dataset_memberships": [], } unique_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) placeholders = ", ".join("?" for _ in unique_document_ids) existing_rows = connection.execute( f""" SELECT id FROM documents WHERE id IN ({placeholders}) """, unique_document_ids, ).fetchall() existing_ids = {int(row["id"]) for row in existing_rows} missing_ids = [document_id for document_id in unique_document_ids if document_id not in existing_ids] if missing_ids: raise RetrieverError(f"Unknown document ids: {', '.join(str(document_id) for document_id in missing_ids)}") current_rows = connection.execute( f""" SELECT DISTINCT document_id FROM dataset_documents WHERE dataset_id = ? AND document_id IN ({placeholders}) """, [dataset_id, *unique_document_ids], ).fetchall() current_ids = {int(row["document_id"]) for row in current_rows} removed_document_ids: list[int] = [] not_present_document_ids: list[int] = [] documents_without_dataset_memberships: list[int] = [] for document_id in unique_document_ids: if document_id not in current_ids: not_present_document_ids.append(document_id) continue connection.execute( """ DELETE FROM dataset_documents WHERE dataset_id = ? AND document_id = ? """, (dataset_id, document_id), ) cached_dataset_id = refresh_document_dataset_cache(connection, document_id) if cached_dataset_id is None: remaining_membership = connection.execute( """ SELECT 1 FROM dataset_documents WHERE document_id = ? LIMIT 1 """, (document_id,), ).fetchone() if remaining_membership is None: documents_without_dataset_memberships.append(document_id) removed_document_ids.append(document_id) if removed_document_ids: connection.execute( """ UPDATE datasets SET updated_at = ? WHERE id = ? """, (utc_now(), dataset_id), ) return { "removed_document_ids": removed_document_ids, "not_present_document_ids": not_present_document_ids, "documents_without_dataset_memberships": documents_without_dataset_memberships, } def delete_dataset_row(connection: sqlite3.Connection, dataset_id: int, root: Path | None = None) -> dict[str, object]: dataset_row = get_dataset_row_by_id(connection, dataset_id) if dataset_row is None: raise RetrieverError(f"Unknown dataset id: {dataset_id}") affected_rows = connection.execute( """ SELECT DISTINCT document_id FROM dataset_documents WHERE dataset_id = ? ORDER BY document_id ASC """, (dataset_id,), ).fetchall() affected_document_ids = [int(row["document_id"]) for row in affected_rows] summary = dataset_summary_by_id(connection, dataset_id, root=root) connection.execute("DELETE FROM datasets WHERE id = ?", (dataset_id,)) documents_without_dataset_memberships: list[int] = [] for document_id in affected_document_ids: row = connection.execute("SELECT id FROM documents WHERE id = ?", (document_id,)).fetchone() if row is None: continue cached_dataset_id = refresh_document_dataset_cache(connection, document_id) if cached_dataset_id is None: remaining_membership = connection.execute( """ SELECT 1 FROM dataset_documents WHERE document_id = ? LIMIT 1 """, (document_id,), ).fetchone() if remaining_membership is None: documents_without_dataset_memberships.append(document_id) return { "deleted_dataset": summary, "affected_document_ids": affected_document_ids, "documents_without_dataset_memberships": documents_without_dataset_memberships, } def prune_unused_filesystem_dataset(connection: sqlite3.Connection) -> bool: dataset_source_row = get_dataset_source_row( connection, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), ) if dataset_source_row is None: return False dataset_id = int(dataset_source_row["dataset_id"]) membership_row = connection.execute( """ SELECT 1 FROM dataset_documents WHERE dataset_id = ? LIMIT 1 """, (dataset_id,), ).fetchone() if membership_row is not None: return False filesystem_document_row = connection.execute( """ SELECT 1 FROM documents WHERE COALESCE(source_kind, ?) = ? LIMIT 1 """, (FILESYSTEM_SOURCE_KIND, FILESYSTEM_SOURCE_KIND), ).fetchone() if filesystem_document_row is not None: return False connection.execute("DELETE FROM datasets WHERE id = ?", (dataset_id,)) return True def container_message_rel_path(source_rel_path: str, source_item_id: str, file_suffix: str) -> str: encoded = encode_source_item_id_for_path(source_item_id) return ( Path(INTERNAL_REL_PATH_PREFIX) / "sources" / Path(source_rel_path) / "messages" / f"{encoded}.{file_suffix}" ).as_posix() def pst_message_rel_path(source_rel_path: str, source_item_id: str) -> str: return container_message_rel_path(source_rel_path, source_item_id, "pstmsg") def mbox_message_rel_path(source_rel_path: str, source_item_id: str) -> str: return container_message_rel_path(source_rel_path, source_item_id, "mboxmsg") def container_preview_file_name(source_item_id: str) -> str: return f"{encode_source_item_id_for_path(source_item_id)}.html" def pst_preview_file_name(source_item_id: str) -> str: return container_preview_file_name(source_item_id) def mbox_preview_file_name(source_item_id: str) -> str: return container_preview_file_name(source_item_id) def container_message_file_name(source_item_id: str, file_suffix: str) -> str: return f"{encode_source_item_id_for_path(source_item_id)}.{file_suffix}" def pst_message_file_name(source_item_id: str) -> str: return container_message_file_name(source_item_id, "pstmsg") def mbox_message_file_name(source_item_id: str) -> str: return container_message_file_name(source_item_id, "mboxmsg") def sanitize_storage_filename(file_name: str) -> str: sanitized = re.sub(r"[\\/:*?\"<>|\x00-\x1f]+", "_", file_name.strip()) sanitized = sanitized.strip().strip(".") return sanitized or "attachment.bin" FILE_TYPE_ALIASES = { "htm": "html", "jpe": "jpg", "jpeg": "jpg", "tiff": "tif", } ATTACHMENT_FILE_TYPE_BY_MIME_TYPE = { "application/json": "json", "application/pdf": "pdf", "application/rtf": "rtf", "application/vnd.ms-excel": "xls", "application/vnd.ms-powerpoint": "ppt", "application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx", "application/xhtml+xml": "html", "application/xml": "xml", "application/zip": "zip", "application/x-zip-compressed": "zip", "text/calendar": "ics", "text/csv": "csv", "text/html": "html", "text/json": "json", "text/markdown": "md", "text/plain": "txt", "text/rtf": "rtf", "text/tab-separated-values": "tsv", "text/tsv": "tsv", "text/xml": "xml", } OLE_COMPOUND_FILE_MAGIC = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" def normalize_file_type_name(raw_value: object) -> str | None: normalized = normalize_whitespace(str(raw_value or "")).lower().lstrip(".") if not normalized: return None return FILE_TYPE_ALIASES.get(normalized, normalized) def normalize_mime_type(raw_value: object) -> str | None: normalized = normalize_whitespace(str(raw_value or "")).lower() if not normalized: return None normalized = normalized.split(";", 1)[0].strip() return normalized or None def attachment_file_type_from_mime_type(raw_value: object) -> str | None: mime_type = normalize_mime_type(raw_value) if mime_type is None or mime_type == "application/octet-stream": return None explicit = ATTACHMENT_FILE_TYPE_BY_MIME_TYPE.get(mime_type) if explicit: return explicit guessed = mimetypes.guess_extension(mime_type, strict=False) return normalize_file_type_name(guessed) def infer_content_type_from_extension(file_type: str) -> str | None: if not file_type: return None if file_type == "md": return "E-Doc" return CONTENT_TYPE_BY_EXTENSION.get(file_type) def canonical_kind_from_metadata( *, extracted_content_type: object = None, extracted_kind: object = None, file_type: object = None, source_kind: object = None, ) -> str: normalized_kind = normalize_whitespace(str(extracted_kind or "")).lower() if normalized_kind in CANONICAL_KIND_VALUES: return normalized_kind normalized_content_type = normalize_whitespace(str(extracted_content_type or "")).lower() normalized_file_type = normalize_whitespace(str(file_type or "")).lower() normalized_source_kind = normalize_whitespace(str(source_kind or "")).lower() if normalized_source_kind in {PST_SOURCE_KIND, MBOX_SOURCE_KIND} or normalized_content_type == "email": return "email" if "spreadsheet" in normalized_content_type or "table" in normalized_content_type: return "spreadsheet" if "presentation" in normalized_content_type: return "presentation" if normalized_content_type == "image": return "image" if normalized_content_type == "source code": return "code" if normalized_content_type == "database": return "data" if normalized_content_type == "container": return "binary" if normalized_content_type in {"calendar", "message"}: return "email" if normalized_content_type in {"chat", "e-doc", "web"}: return "document" if normalized_file_type in {"csv", "json", "tsv", "xml", "yaml", "yml"}: return "data" if normalized_file_type in {"xls", "xlsx", "xlsm", "xlsb", "ods", "numbers"}: return "spreadsheet" if normalized_file_type in {"ppt", "pptx", "pptm", "odp", "key"}: return "presentation" if normalized_file_type in {"png", "jpg", "jpeg", "gif", "bmp", "tif", "tiff", "webp", "svg"}: return "image" if normalized_file_type in CURATED_TEXT_SOURCE_FILE_TYPES: return "code" if normalized_file_type in {"doc", "docx", "pdf", "txt", "rtf", "md", "html", "htm"}: return "document" return "unknown" def canonical_kind_compatible(left: object, right: object) -> bool: left_kind = normalize_whitespace(str(left or "")).lower() or "unknown" right_kind = normalize_whitespace(str(right or "")).lower() or "unknown" return left_kind == "unknown" or right_kind == "unknown" or left_kind == right_kind def occurrence_field_count(row: sqlite3.Row | dict[str, object]) -> int: fields = [ "extracted_author", "extracted_title", "extracted_subject", "extracted_participants", "extracted_recipients", "extracted_doc_authored_at", "extracted_doc_modified_at", "extracted_content_type", "extracted_kind", ] return sum(1 for field_name in fields if normalize_whitespace(str(row[field_name] or ""))) def text_status_priority(status: object) -> int: normalized = normalize_whitespace(str(status or "")).lower() return TEXT_STATUS_PRIORITIES.get(normalized, max(TEXT_STATUS_PRIORITIES.values()) + 1) def source_kind_priority(source_kind: object) -> int: normalized = normalize_whitespace(str(source_kind or "")).lower() return SOURCE_KIND_PREFERRED_ORDER.get(normalized, max(SOURCE_KIND_PREFERRED_ORDER.values()) + 1) def active_occurrence_rows_for_document( connection: sqlite3.Connection, document_id: int, *, include_all_statuses: bool = False, ) -> list[sqlite3.Row]: if include_all_statuses: rows = connection.execute( """ SELECT * FROM document_occurrences WHERE document_id = ? ORDER BY id ASC """, (document_id,), ).fetchall() else: rows = connection.execute( """ SELECT * FROM document_occurrences WHERE document_id = ? AND lifecycle_status = ? ORDER BY id ASC """, (document_id, ACTIVE_OCCURRENCE_STATUS), ).fetchall() return rows def select_preferred_occurrence(rows: list[sqlite3.Row]) -> sqlite3.Row | None: if not rows: return None ranked_rows = sorted( rows, key=lambda row: ( 0 if row["lifecycle_status"] == ACTIVE_OCCURRENCE_STATUS else 1, source_kind_priority(row["source_kind"]), text_status_priority(row["text_status"]), -occurrence_field_count(row), 0 if int(row["has_preview"] or 0) else 1, parse_utc_timestamp(row["ingested_at"]) or datetime.max.replace(tzinfo=timezone.utc), int(row["id"]), ), ) return ranked_rows[0] def occurrence_field_value( preferred_row: sqlite3.Row | None, active_rows: list[sqlite3.Row], preferred_column: str, ) -> object: if preferred_row is not None: preferred_value = preferred_row[preferred_column] if preferred_value not in (None, ""): return preferred_value sorted_rows = sorted( active_rows, key=lambda row: ( parse_utc_timestamp(row["ingested_at"]) or datetime.max.replace(tzinfo=timezone.utc), int(row["id"]), ), ) for row in sorted_rows: value = row[preferred_column] if value not in (None, ""): return value return None def normalize_custodian_values(values: list[object] | tuple[object, ...] | set[object]) -> list[str]: normalized_values: list[str] = [] seen: set[str] = set() for raw_value in values: normalized = normalize_whitespace(str(raw_value or "")) if not normalized or normalized in seen: continue seen.add(normalized) normalized_values.append(normalized) return normalized_values def occurrence_rows_in_preferred_order(rows: list[sqlite3.Row]) -> list[sqlite3.Row]: return sorted( rows, key=lambda row: ( 0 if row["lifecycle_status"] == ACTIVE_OCCURRENCE_STATUS else 1, source_kind_priority(row["source_kind"]), text_status_priority(row["text_status"]), -occurrence_field_count(row), 0 if int(row["has_preview"] or 0) else 1, parse_utc_timestamp(row["ingested_at"]) or datetime.max.replace(tzinfo=timezone.utc), int(row["id"]), ), ) def custodian_values_from_occurrence_rows(rows: list[sqlite3.Row]) -> list[str]: return normalize_custodian_values([row["custodian"] for row in occurrence_rows_in_preferred_order(rows)]) def parse_document_custodians_json(raw_value: object) -> list[str]: if isinstance(raw_value, list): return normalize_custodian_values(list(raw_value)) if not raw_value: return [] try: parsed = json.loads(str(raw_value)) except (TypeError, ValueError, json.JSONDecodeError): return normalize_custodian_values([raw_value]) if isinstance(parsed, list): return normalize_custodian_values(parsed) return normalize_custodian_values([parsed]) def document_custodian_values_from_row(row: sqlite3.Row | dict[str, object] | None) -> list[str]: if row is None: return [] if "custodians_json" in row.keys(): # type: ignore[attr-defined] return parse_document_custodians_json(row["custodians_json"]) # type: ignore[index] if "custodian" in row.keys(): # type: ignore[attr-defined] return normalize_custodian_values([row["custodian"]]) # type: ignore[index] return [] def document_custodian_display_text_from_row(row: sqlite3.Row | dict[str, object] | None) -> str | None: values = document_custodian_values_from_row(row) if not values: return None return ", ".join(values) def normalize_entity_text(value: object) -> str: return normalize_whitespace(unicodedata.normalize("NFKC", str(value or ""))) def normalize_entity_lookup_text(value: object) -> str: text = normalize_entity_text(value).lower() text = re.sub(r"[^\w@.+#/-]+", " ", text, flags=re.UNICODE) return normalize_whitespace(text) def normalize_entity_email(value: object) -> str | None: text = normalize_entity_text(value).strip("<>;,") if not text or "@" not in text: return None candidate = text.lower() candidate = re.sub(r"^mailto:", "", candidate) candidate = candidate.strip("<>;,") if not re.fullmatch(r"[a-z0-9.!#$%&'*+/=?^_`{|}~-]+@[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?", candidate): return None if "." not in candidate.rsplit("@", 1)[1]: return None return candidate def normalize_entity_identifier_name(value: object) -> str | None: text = normalize_entity_lookup_text(value) text = re.sub(r"[^a-z0-9]+", "_", text).strip("_") return text or None def normalize_entity_handle(value: object) -> str | None: text = normalize_entity_text(value).strip() if not text: return None if text.startswith("@"): text = text[1:] text = text.strip() if not text: return None normalized = re.sub(r"\s+", "", text).lower() if not re.fullmatch(r"[a-z0-9._-]{2,128}", normalized): return None return normalized def normalize_entity_phone(value: object) -> dict[str, object] | None: text = normalize_entity_text(value) if not text: return None extension = None extension_match = re.search(r"(?i)(?:ext\.?|extension|x)\s*([0-9]{1,10})\b", text) phone_text = text if extension_match: extension = extension_match.group(1) phone_text = (text[:extension_match.start()] + text[extension_match.end():]).strip() digits = re.sub(r"\D+", "", phone_text) if len(digits) < 7 or len(digits) > 15: return None has_phone_shape = bool(re.search(r"[()+\-\s.]", phone_text)) or phone_text.strip().startswith("+") if not has_phone_shape and len(digits) < 10: return None if phone_text.strip().startswith("+"): base_phone = f"+{digits}" elif len(digits) == 10: base_phone = f"+1{digits}" elif len(digits) == 11 and digits.startswith("1"): base_phone = f"+{digits}" else: base_phone = digits normalized = f"{base_phone}x{extension}" if extension else base_phone return { "display_value": text, "normalized_value": normalized, "parsed_phone": { "base_phone": base_phone, "extension": extension, }, } def strip_entity_container_words(value: str) -> str: text = normalize_entity_text(value) stripped = normalize_whitespace(re.sub(r"(?i)\b(?:mailbox|archive|export|pst|mbox)\b", " ", text)) if len(entity_name_tokens(stripped)) >= 2: return stripped return text def entity_name_tokens(value: object) -> list[str]: text = normalize_entity_text(value) text = re.sub(r"[\"'()<>]", " ", text) text = re.sub(r"\b(?:mr|mrs|ms|miss|dr|prof)\.?\b", " ", text, flags=re.IGNORECASE) text = re.sub(r"\b(?:jr|sr|ii|iii|iv|esq)\.?\b", " ", text, flags=re.IGNORECASE) return [ token for token in re.split(r"[\s.]+", text) if token and re.search(r"[A-Za-z]", token) ] def parse_entity_name(value: object) -> dict[str, object] | None: text = strip_entity_container_words(str(value or "")) if not text or "@" in text: return None compact = normalize_entity_lookup_text(text) if not compact: return None parts = [normalize_entity_text(part) for part in text.split(",", 1)] parsed: dict[str, object] = {} display_name = normalize_entity_text(text) if len(parts) == 2 and parts[0] and parts[1]: family_tokens = entity_name_tokens(parts[0]) given_tokens = entity_name_tokens(parts[1]) if family_tokens and given_tokens: tokens = [*given_tokens, *family_tokens] display_name = " ".join(tokens) parsed = { "given": given_tokens[0], "middle": " ".join(given_tokens[1:]) or None, "family": " ".join(family_tokens), } if not parsed: tokens = entity_name_tokens(text) if not tokens: return None if len(tokens) >= 2: parsed = { "given": tokens[0], "middle": " ".join(tokens[1:-1]) or None, "family": tokens[-1], } display_name = " ".join(tokens) else: parsed = { "given": tokens[0], "middle": None, "family": None, } display_name = tokens[0] display_tokens = entity_name_tokens(display_name) if not display_tokens: return None normalized_full_name = normalize_entity_lookup_text(" ".join(display_tokens)) family = normalize_entity_lookup_text(str(parsed.get("family") or "")) given = normalize_entity_lookup_text(str(parsed.get("given") or "")) middle = normalize_entity_lookup_text(str(parsed.get("middle") or "")) sort_parts = [part for part in (family, given, middle) if part] normalized_sort_name = normalize_entity_lookup_text(" ".join(sort_parts)) if sort_parts else normalized_full_name return { "display_value": display_name, "normalized_value": normalized_full_name, "parsed_name": {key: value for key, value in parsed.items() if value}, "normalized_full_name": normalized_full_name, "normalized_sort_name": normalized_sort_name, "is_full_name": len(display_tokens) >= 2 and bool(parsed.get("family")), } def entity_name_identifier_looks_like_export_artifact(raw_value: object) -> bool: display_value = normalize_entity_text(raw_value) normalized = normalize_entity_lookup_text(display_value) if not normalized: return False compact = re.sub(r"\s+", "", normalized) if compact.startswith("-"): return True letters = re.sub(r"[^a-z]", "", compact.lower()) if letters and not re.search(r"[aeiou]", letters): return True if re.fullmatch(r"[A-Za-z0-9_-]{4,8}", display_value): has_upper = bool(re.search(r"[A-Z]", display_value)) has_lower = bool(re.search(r"[a-z]", display_value)) has_digit_or_separator = bool(re.search(r"[0-9_-]", display_value)) if has_upper and has_lower and has_digit_or_separator: return True return False def entity_type_from_candidate_parts( *, name_value: str | None, email_value: str | None, name_is_full: bool = False, ) -> str: email = normalize_entity_email(email_value) if email: local_part = email.split("@", 1)[0] normalized_local = re.sub(r"[^a-z0-9]+", "", local_part.lower()) if normalized_local in SYSTEM_MAILBOX_LOCAL_PARTS or any( token in normalized_local for token in SYSTEM_MAILBOX_LOCAL_CONTAINS ): return ENTITY_TYPE_SYSTEM_MAILBOX if normalized_local in SHARED_MAILBOX_LOCAL_PARTS: return ENTITY_TYPE_SHARED_MAILBOX name = normalize_entity_lookup_text(name_value or "") if email and any(re.search(rf"\b{re.escape(hint)}\b", name) for hint in SHARED_MAILBOX_NAME_HINTS): return ENTITY_TYPE_SHARED_MAILBOX if re.search(r"\b(inc|llc|llp|ltd|corp|corporation|company|co|plc|gmbh|sarl|partners|holdings)\b", name): return ENTITY_TYPE_ORGANIZATION if email or (name and name_is_full): return ENTITY_TYPE_PERSON if name: return ENTITY_TYPE_UNKNOWN return ENTITY_TYPE_UNKNOWN def entity_candidate_identifier_key(identifier: dict[str, object]) -> str: identifier_type = str(identifier.get("identifier_type") or "") pieces = [identifier_type] for field_name in ("provider", "provider_scope", "identifier_name", "identifier_scope", "normalized_value"): value = normalize_entity_lookup_text(identifier.get(field_name) or "") if value: pieces.append(value) return ":".join(pieces) def entity_candidate_key(role: str, identifiers: list[dict[str, object]], fallback_value: object) -> str: identifier_keys = sorted(entity_candidate_identifier_key(identifier) for identifier in identifiers) if identifier_keys: return "|".join([role, *identifier_keys]) return "|".join([role, "raw", normalize_entity_lookup_text(fallback_value)]) def split_entity_like_values(raw_value: object, *, prefer_single_comma_name: bool = False) -> list[str]: text = normalize_entity_text(raw_value) if not text: return [] semicolon_parts = [normalize_entity_text(part) for part in text.split(";")] if len([part for part in semicolon_parts if part]) > 1: return [part for part in semicolon_parts if part] if "@" in text or "<" in text: parsed_addresses = getaddresses([text]) parts: list[str] = [] for display_name, address in parsed_addresses: normalized_email = normalize_entity_email(address) display_name = normalize_entity_text(display_name) if normalized_email: parts.append(f"{display_name} <{normalized_email}>" if display_name else normalized_email) elif display_name: parts.append(display_name) if parts: return parts comma_count = text.count(",") if prefer_single_comma_name and comma_count == 1 and parse_entity_name(text): return [text] if comma_count: return [part for part in (normalize_entity_text(part) for part in text.split(",")) if part] return [text] ENTITY_LABELED_EXTERNAL_ID_PATTERN = re.compile( r"\b([A-Za-z][A-Za-z0-9 _/-]{1,40})\s*[:=#]\s*([A-Za-z0-9][A-Za-z0-9_.:/-]{2,80})" ) def parse_labeled_external_ids(value: object) -> list[dict[str, object]]: text = normalize_entity_text(value) identifiers: list[dict[str, object]] = [] for match in ENTITY_LABELED_EXTERNAL_ID_PATTERN.finditer(text): identifier_name = normalize_entity_identifier_name(match.group(1)) normalized_value = normalize_entity_lookup_text(match.group(2)) if not identifier_name or not normalized_value: continue identifiers.append( { "identifier_type": "external_id", "display_value": match.group(2), "normalized_value": normalized_value, "identifier_name": identifier_name, } ) return identifiers def strip_labeled_external_ids(value: object) -> str: text = normalize_entity_text(value) if not text: return "" stripped = ENTITY_LABELED_EXTERNAL_ID_PATTERN.sub(" ", text) stripped = re.sub(r"\s*[\[\](){}]\s*", " ", stripped) return normalize_entity_text(stripped) def parse_entity_candidate_text( raw_value: object, *, role: str, provider: str | None = None, provider_scope: str | None = None, ) -> dict[str, object] | None: text = normalize_entity_text(raw_value) if not text: return None emails: list[str] = [] names: list[str] = [] angle_address_match = re.match(r"^(?P<display>.*?)<(?P<address>[^>]+)>\s*$", text) if angle_address_match: email = normalize_entity_email(angle_address_match.group("address")) display_name = normalize_entity_text(angle_address_match.group("display")) if email: emails.append(email) if display_name: names.append(display_name.strip("\"' ")) if not emails: parsed_addresses = getaddresses([text]) for display_name, address in parsed_addresses: email = normalize_entity_email(address) if email and email not in emails: emails.append(email) if normalize_entity_text(display_name): names.append(normalize_entity_text(display_name)) if not emails: for email_match in re.finditer(r"[A-Za-z0-9.!#$%&'*+/=?^_`{|}~-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}", text): email = normalize_entity_email(email_match.group(0)) if email and email not in emails: emails.append(email) name_source = names[0] if names else re.sub(r"<[^>]+>", " ", text) for email in emails: name_source = re.sub(re.escape(email), " ", name_source, flags=re.IGNORECASE) phone = normalize_entity_phone(text) if phone: name_source = normalize_whitespace(str(name_source).replace(str(phone["display_value"]), " ")) name_source = strip_labeled_external_ids(name_source) identifiers: list[dict[str, object]] = [] for email in emails: identifiers.append( { "identifier_type": "email", "display_value": email, "normalized_value": email, "is_verified": 1, } ) parsed_name = parse_entity_name(name_source) if parsed_name is not None and entity_name_identifier_looks_like_export_artifact(parsed_name["display_value"]): parsed_name = None if parsed_name is not None: identifiers.append( { "identifier_type": "name", "display_value": parsed_name["display_value"], "normalized_value": parsed_name["normalized_value"], "parsed_name_json": json.dumps(parsed_name["parsed_name"], ensure_ascii=True, sort_keys=True), "normalized_full_name": parsed_name["normalized_full_name"], "normalized_sort_name": parsed_name["normalized_sort_name"], "is_primary": 1 if parsed_name["is_full_name"] else 0, } ) if phone is not None: identifiers.append( { "identifier_type": "phone", "display_value": phone["display_value"], "normalized_value": phone["normalized_value"], "parsed_phone_json": json.dumps(phone["parsed_phone"], ensure_ascii=True, sort_keys=True), } ) if provider and provider_scope: handle = normalize_entity_handle(text) if handle: identifiers.append( { "identifier_type": "handle", "display_value": text, "normalized_value": handle, "provider": normalize_entity_identifier_name(provider), "provider_scope": normalize_entity_lookup_text(provider_scope), } ) identifiers.extend(parse_labeled_external_ids(text)) if not identifiers: return None display_basis = ( str(parsed_name["display_value"]) if parsed_name is not None else emails[0] if emails else str(phone["display_value"]) if phone is not None else text ) entity_type = entity_type_from_candidate_parts( name_value=str(parsed_name["display_value"]) if parsed_name is not None else None, email_value=emails[0] if emails else None, name_is_full=bool(parsed_name.get("is_full_name")) if parsed_name is not None else False, ) return { "role": role, "raw_value": text, "display_value": display_basis, "entity_type": entity_type, "identifiers": identifiers, "normalized_candidate_key": entity_candidate_key(role, identifiers, text), } def normalize_entity_hint_identifier( raw_identifier: object, *, default_source_kind: object = None, ) -> dict[str, object] | None: if not isinstance(raw_identifier, dict): return None identifier_type = normalize_entity_identifier_name( raw_identifier.get("identifier_type") or raw_identifier.get("type") or "" ) if identifier_type not in {"email", "handle", "name", "external_id"}: return None display_value = normalize_entity_text( raw_identifier.get("display_value") or raw_identifier.get("value") or raw_identifier.get("normalized_value") or "" ) is_primary = 1 if int(raw_identifier.get("is_primary") or 0) else 0 is_verified = 1 if int(raw_identifier.get("is_verified") or 0) else 0 source_kind = normalize_entity_identifier_name(raw_identifier.get("source_kind") or default_source_kind or "") identifier: dict[str, object] if identifier_type == "email": email = normalize_entity_email(raw_identifier.get("normalized_value") or display_value) if not email: return None identifier = { "identifier_type": "email", "display_value": email, "normalized_value": email, "is_primary": is_primary, "is_verified": is_verified, } elif identifier_type == "handle": provider = normalize_entity_identifier_name(raw_identifier.get("provider") or "") provider_scope = normalize_entity_lookup_text(raw_identifier.get("provider_scope") or raw_identifier.get("scope") or "") handle = normalize_entity_handle(raw_identifier.get("normalized_value") or display_value) if not provider or not provider_scope or not handle: return None identifier = { "identifier_type": "handle", "display_value": display_value or f"@{handle}", "normalized_value": handle, "provider": provider, "provider_scope": provider_scope, "is_primary": is_primary, "is_verified": is_verified, } elif identifier_type == "name": parsed_name = parse_entity_name(raw_identifier.get("normalized_value") or display_value) if parsed_name is None: return None identifier = { "identifier_type": "name", "display_value": parsed_name["display_value"], "normalized_value": parsed_name["normalized_value"], "parsed_name_json": json.dumps(parsed_name["parsed_name"], ensure_ascii=True, sort_keys=True), "normalized_full_name": parsed_name["normalized_full_name"], "normalized_sort_name": parsed_name["normalized_sort_name"], "is_primary": is_primary or (1 if parsed_name["is_full_name"] else 0), "is_verified": is_verified, } else: identifier_name = normalize_entity_identifier_name( raw_identifier.get("identifier_name") or raw_identifier.get("name") or "" ) normalized_value = normalize_entity_lookup_text(raw_identifier.get("normalized_value") or display_value) if not identifier_name or not normalized_value: return None identifier = { "identifier_type": "external_id", "display_value": display_value or normalized_value, "normalized_value": normalized_value, "identifier_name": identifier_name, "is_primary": is_primary, "is_verified": is_verified, } identifier_scope = normalize_entity_lookup_text( raw_identifier.get("identifier_scope") or raw_identifier.get("scope") or "" ) if identifier_scope: identifier["identifier_scope"] = identifier_scope if source_kind: identifier["source_kind"] = source_kind return identifier def parse_entity_hint_candidates( raw_hints: object, *, role: str, ) -> list[dict[str, object]]: if not isinstance(raw_hints, list): return [] candidates: list[dict[str, object]] = [] seen_keys: set[str] = set() for raw_hint in raw_hints: if not isinstance(raw_hint, dict): continue display_value = normalize_entity_text( raw_hint.get("display_value") or raw_hint.get("name") or raw_hint.get("value") or "" ) if not display_value: continue base_candidate = parse_entity_candidate_text(display_value, role=role) identifiers = list(base_candidate.get("identifiers") or []) if base_candidate is not None else [] seen_identifier_keys = { entity_candidate_identifier_key(identifier) for identifier in identifiers } for raw_identifier in list(raw_hint.get("identifiers") or []): identifier = normalize_entity_hint_identifier( raw_identifier, default_source_kind=raw_hint.get("source_kind"), ) if identifier is None: continue identifier_key = entity_candidate_identifier_key(identifier) if identifier_key in seen_identifier_keys: continue seen_identifier_keys.add(identifier_key) identifiers.append(identifier) if not identifiers: continue parsed_name_identifiers = [identifier for identifier in identifiers if identifier.get("identifier_type") == "name"] email_identifiers = [identifier for identifier in identifiers if identifier.get("identifier_type") == "email"] derived_entity_type = entity_type_from_candidate_parts( name_value=str(parsed_name_identifiers[0]["display_value"]) if parsed_name_identifiers else display_value, email_value=str(email_identifiers[0]["normalized_value"]) if email_identifiers else None, name_is_full=bool(parsed_name_identifiers and int(parsed_name_identifiers[0].get("is_primary") or 0)), ) entity_type = str(base_candidate.get("entity_type")) if base_candidate is not None else derived_entity_type if entity_type == ENTITY_TYPE_UNKNOWN and derived_entity_type != ENTITY_TYPE_UNKNOWN: entity_type = derived_entity_type display_basis = str(base_candidate.get("display_value")) if base_candidate is not None else display_value candidate = { "role": role, "raw_value": display_value, "display_value": display_basis, "entity_type": entity_type, "identifiers": identifiers, "normalized_candidate_key": entity_candidate_key(role, identifiers, display_value), } key = str(candidate["normalized_candidate_key"]) if key in seen_keys: continue seen_keys.add(key) candidates.append(candidate) return candidates def entity_candidate_match_keys(candidate: dict[str, object]) -> set[str]: keys = { normalize_entity_lookup_text(candidate.get("display_value") or ""), normalize_entity_lookup_text(candidate.get("raw_value") or ""), } for identifier in list(candidate.get("identifiers") or []): if not isinstance(identifier, dict): continue if identifier.get("identifier_type") != "name": continue for field_name in ("display_value", "normalized_value", "normalized_full_name", "normalized_sort_name"): key = normalize_entity_lookup_text(identifier.get(field_name) or "") if key: keys.add(key) return {key for key in keys if key} def entity_hints_for_role(raw_hints: object, role: str) -> list[dict[str, object]]: if isinstance(raw_hints, str): try: raw_hints = json.loads(raw_hints) except (TypeError, ValueError, json.JSONDecodeError): return [] if not isinstance(raw_hints, dict): return [] role_keys = [role] if role == "participant": role_keys.append("participants") elif role == "recipient": role_keys.append("recipients") elif role == "custodian": role_keys.append("custodians") for role_key in role_keys: value = raw_hints.get(role_key) if isinstance(value, list): return [item for item in value if isinstance(item, dict)] return [] def parse_entity_candidates_with_hints( raw_value: object, *, role: str, raw_hints: object = None, ) -> list[dict[str, object]]: hint_candidates = parse_entity_hint_candidates(entity_hints_for_role(raw_hints, role), role=role) covered_keys: set[str] = set() for candidate in hint_candidates: covered_keys.update(entity_candidate_match_keys(candidate)) candidates = list(hint_candidates) seen_keys = {str(candidate["normalized_candidate_key"]) for candidate in candidates} for candidate in parse_entity_candidates(raw_value, role=role): if entity_candidate_match_keys(candidate) & covered_keys: continue key = str(candidate["normalized_candidate_key"]) if key in seen_keys: continue seen_keys.add(key) candidates.append(candidate) return candidates def parse_entity_candidates( raw_value: object, *, role: str, provider: str | None = None, provider_scope: str | None = None, ) -> list[dict[str, object]]: candidates: list[dict[str, object]] = [] seen_keys: set[str] = set() for part in split_entity_like_values(raw_value, prefer_single_comma_name=role in {"author", "custodian"}): candidate = parse_entity_candidate_text(part, role=role, provider=provider, provider_scope=provider_scope) if candidate is None: continue key = str(candidate["normalized_candidate_key"]) if key in seen_keys: continue seen_keys.add(key) candidates.append(candidate) return candidates def source_backed_dataset_policy_from_row(row: sqlite3.Row | None) -> dict[str, object]: if row is None or normalize_whitespace(str(row["source_kind"] or "")).lower() == MANUAL_DATASET_SOURCE_KIND: return { "dataset_id": None, "allow_auto_merge": False, "email_auto_merge": False, "handle_auto_merge": False, "phone_auto_merge": False, "name_auto_merge": False, "external_id_auto_merge_names": set(), } names = normalize_string_list(row["external_id_auto_merge_names_json"]) return { "dataset_id": int(row["id"]), "allow_auto_merge": bool(int(row["allow_auto_merge"] or 0)), "email_auto_merge": bool(int(row["email_auto_merge"] or 0)), "handle_auto_merge": bool(int(row["handle_auto_merge"] or 0)), "phone_auto_merge": bool(int(row["phone_auto_merge"] or 0)), "name_auto_merge": bool(int(row["name_auto_merge"] or 0)), "external_id_auto_merge_names": { name for raw_name in names for name in [normalize_entity_identifier_name(raw_name)] if name }, } def dataset_merge_policy_payload_from_row(row: sqlite3.Row) -> dict[str, object]: source_kind = normalize_whitespace(str(row["source_kind"] or "")).lower() names = sorted( { name for raw_name in normalize_string_list(row["external_id_auto_merge_names_json"]) for name in [normalize_entity_identifier_name(raw_name)] if name } ) return { "dataset_id": int(row["id"]), "source_backed": source_kind != MANUAL_DATASET_SOURCE_KIND, "allow_auto_merge": bool(int(row["allow_auto_merge"] or 0)), "email_auto_merge": bool(int(row["email_auto_merge"] or 0)), "handle_auto_merge": bool(int(row["handle_auto_merge"] or 0)), "phone_auto_merge": bool(int(row["phone_auto_merge"] or 0)), "name_auto_merge": bool(int(row["name_auto_merge"] or 0)), "external_id_auto_merge_names": names, } def source_backed_dataset_policy_for_source( connection: sqlite3.Connection, dataset_source_row: sqlite3.Row | None, ) -> dict[str, object]: if dataset_source_row is None: return source_backed_dataset_policy_from_row(None) dataset_row = get_dataset_row_by_id(connection, int(dataset_source_row["dataset_id"])) return source_backed_dataset_policy_from_row(dataset_row) def identifier_auto_merge_enabled(identifier: dict[str, object], policy: dict[str, object]) -> bool: if not bool(policy.get("allow_auto_merge")): return False identifier_type = str(identifier.get("identifier_type") or "") if identifier_type == "email": return bool(policy.get("email_auto_merge")) if identifier_type == "handle": return bool(policy.get("handle_auto_merge")) and bool(identifier.get("provider")) and bool(identifier.get("provider_scope")) if identifier_type == "phone": return bool(policy.get("phone_auto_merge")) if identifier_type == "name": return bool(policy.get("name_auto_merge")) and bool(identifier.get("normalized_full_name")) and bool(identifier.get("normalized_sort_name")) if identifier_type == "external_id": enabled_names = policy.get("external_id_auto_merge_names") return isinstance(enabled_names, set) and str(identifier.get("identifier_name") or "") in enabled_names return False def resolution_key_lookup_clause(identifier: dict[str, object]) -> tuple[str, list[object]]: identifier_type = str(identifier.get("identifier_type") or "") normalized_value = str(identifier.get("normalized_value") or "") if identifier_type == "handle": return ( """ key_type = ? AND provider = ? AND provider_scope = ? AND normalized_value = ? """, [ identifier_type, identifier.get("provider"), identifier.get("provider_scope"), normalized_value, ], ) if identifier_type == "external_id": return ( """ key_type = ? AND identifier_name = ? AND COALESCE(identifier_scope, '') = COALESCE(?, '') AND normalized_value = ? """, [ identifier_type, identifier.get("identifier_name"), identifier.get("identifier_scope"), normalized_value, ], ) return "key_type = ? AND normalized_value = ?", [identifier_type, normalized_value] def canonicalize_entity_id(connection: sqlite3.Connection, entity_id: int) -> int | None: seen: set[int] = set() current_id = int(entity_id) while current_id not in seen: seen.add(current_id) row = connection.execute( """ SELECT id, canonical_status, merged_into_entity_id FROM entities WHERE id = ? """, (current_id,), ).fetchone() if row is None: return None if row["canonical_status"] != ENTITY_STATUS_MERGED: return int(row["id"]) if row["merged_into_entity_id"] is None: return None current_id = int(row["merged_into_entity_id"]) return None def entity_types_conflict(left_type: object, right_type: object) -> bool: left = normalize_entity_lookup_text(left_type) right = normalize_entity_lookup_text(right_type) if not left or not right or left == right: return False if ENTITY_TYPE_UNKNOWN in {left, right}: return False concrete = {ENTITY_TYPE_PERSON, ENTITY_TYPE_ORGANIZATION, ENTITY_TYPE_SHARED_MAILBOX, ENTITY_TYPE_SYSTEM_MAILBOX} return left in concrete and right in concrete def active_entity_id_for_resolution_key( connection: sqlite3.Connection, identifier: dict[str, object], ) -> int | None: clause, params = resolution_key_lookup_clause(identifier) row = connection.execute( f""" SELECT entity_id FROM entity_resolution_keys WHERE {clause} ORDER BY id ASC LIMIT 1 """, params, ).fetchone() if row is None: return None entity_id = canonicalize_entity_id(connection, int(row["entity_id"])) if entity_id is None: return None entity_row = connection.execute( """ SELECT canonical_status FROM entities WHERE id = ? """, (entity_id,), ).fetchone() if entity_row is None or entity_row["canonical_status"] != ENTITY_STATUS_ACTIVE: return None return entity_id def create_entity_for_candidate(connection: sqlite3.Connection, candidate: dict[str, object]) -> int: now = utc_now() connection.execute( """ INSERT INTO entities ( entity_type, display_name, display_name_source, entity_origin, canonical_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( candidate.get("entity_type") or ENTITY_TYPE_UNKNOWN, None, ENTITY_DISPLAY_SOURCE_AUTO, ENTITY_ORIGIN_OBSERVED, ENTITY_STATUS_ACTIVE, now, now, ), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) def ensure_entity_identifier( connection: sqlite3.Connection, *, entity_id: int, identifier: dict[str, object], ) -> int: identifier_type = str(identifier.get("identifier_type") or "") normalized_value = str(identifier.get("normalized_value") or "") if not identifier_type or not normalized_value: raise RetrieverError("Entity identifiers require type and normalized value.") provider = identifier.get("provider") provider_scope = identifier.get("provider_scope") identifier_name = identifier.get("identifier_name") identifier_scope = identifier.get("identifier_scope") existing_row = connection.execute( """ SELECT id FROM entity_identifiers WHERE entity_id = ? AND identifier_type = ? AND normalized_value = ? AND COALESCE(provider, '') = COALESCE(?, '') AND COALESCE(provider_scope, '') = COALESCE(?, '') AND COALESCE(identifier_name, '') = COALESCE(?, '') AND COALESCE(identifier_scope, '') = COALESCE(?, '') ORDER BY id ASC LIMIT 1 """, (entity_id, identifier_type, normalized_value, provider, provider_scope, identifier_name, identifier_scope), ).fetchone() now = utc_now() if existing_row is not None: connection.execute( """ UPDATE entity_identifiers SET display_value = COALESCE(NULLIF(display_value, ''), ?), parsed_name_json = COALESCE(parsed_name_json, ?), parsed_phone_json = COALESCE(parsed_phone_json, ?), normalized_full_name = COALESCE(normalized_full_name, ?), normalized_sort_name = COALESCE(normalized_sort_name, ?), is_primary = CASE WHEN ? THEN 1 ELSE is_primary END, is_verified = CASE WHEN ? THEN 1 ELSE is_verified END, updated_at = ? WHERE id = ? """, ( identifier.get("display_value") or normalized_value, identifier.get("parsed_name_json"), identifier.get("parsed_phone_json"), identifier.get("normalized_full_name"), identifier.get("normalized_sort_name"), 1 if int(identifier.get("is_primary") or 0) else 0, 1 if int(identifier.get("is_verified") or 0) else 0, now, int(existing_row["id"]), ), ) return int(existing_row["id"]) connection.execute( """ INSERT INTO entity_identifiers ( entity_id, identifier_type, display_value, normalized_value, provider, provider_scope, identifier_name, identifier_scope, parsed_name_json, parsed_phone_json, normalized_full_name, normalized_sort_name, is_primary, is_verified, source_kind, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( entity_id, identifier_type, str(identifier.get("display_value") or normalized_value), normalized_value, provider, provider_scope, identifier_name, identifier_scope, identifier.get("parsed_name_json"), identifier.get("parsed_phone_json"), identifier.get("normalized_full_name"), identifier.get("normalized_sort_name"), 1 if int(identifier.get("is_primary") or 0) else 0, 1 if int(identifier.get("is_verified") or 0) else 0, str(identifier.get("source_kind") or "auto"), now, now, ), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) def ensure_entity_resolution_key( connection: sqlite3.Connection, *, entity_id: int, identifier_id: int, identifier: dict[str, object], ) -> int | None: now = utc_now() try: connection.execute( """ INSERT INTO entity_resolution_keys ( entity_id, identifier_id, key_type, provider, provider_scope, identifier_name, identifier_scope, normalized_value, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( entity_id, identifier_id, identifier.get("identifier_type"), identifier.get("provider"), identifier.get("provider_scope"), identifier.get("identifier_name"), identifier.get("identifier_scope"), identifier.get("normalized_value"), now, now, ), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) except sqlite3.IntegrityError: existing_owner = active_entity_id_for_resolution_key(connection, identifier) if existing_owner == entity_id: return None return None def recompute_entity_caches(connection: sqlite3.Connection, entity_id: int) -> None: row = connection.execute( """ SELECT * FROM entities WHERE id = ? """, (entity_id,), ).fetchone() if row is None or row["canonical_status"] != ENTITY_STATUS_ACTIVE: return identifiers = connection.execute( """ SELECT * FROM entity_identifiers WHERE entity_id = ? ORDER BY is_primary DESC, is_verified DESC, id ASC """, (entity_id,), ).fetchall() emails = [item for item in identifiers if item["identifier_type"] == "email"] phones = [item for item in identifiers if item["identifier_type"] == "phone"] names = [item for item in identifiers if item["identifier_type"] == "name"] primary_email = str(emails[0]["normalized_value"]) if emails else None primary_phone = str(phones[0]["normalized_value"]) if phones else None display_name = row["display_name"] sort_name = row["sort_name"] if row["display_name_source"] != ENTITY_DISPLAY_SOURCE_MANUAL: full_name_rows = [ item for item in names if normalize_whitespace(str(item["normalized_full_name"] or "")) and len(str(item["normalized_full_name"]).split()) >= 2 ] display_name_rows = full_name_rows if row["entity_type"] != ENTITY_TYPE_PERSON and not display_name_rows: display_name_rows = names if display_name_rows: display_name = str(display_name_rows[0]["display_value"]) sort_name = str(display_name_rows[0]["normalized_sort_name"] or display_name_rows[0]["normalized_full_name"]) else: display_name = None sort_name = None resolution_key_row = connection.execute( """ SELECT 1 FROM entity_resolution_keys WHERE entity_id = ? LIMIT 1 """, (entity_id,), ).fetchone() entity_origin = row["entity_origin"] if entity_origin != ENTITY_ORIGIN_MANUAL: entity_origin = ENTITY_ORIGIN_IDENTIFIED if resolution_key_row is not None else ENTITY_ORIGIN_OBSERVED connection.execute( """ UPDATE entities SET display_name = ?, primary_email = ?, primary_phone = ?, sort_name = ?, entity_origin = ?, updated_at = ? WHERE id = ? """, (display_name, primary_email, primary_phone, sort_name, entity_origin, utc_now(), entity_id), ) def resolve_entity_candidate( connection: sqlite3.Connection, candidate: dict[str, object], *, policy: dict[str, object], ) -> int: identifiers = list(candidate.get("identifiers") or []) enabled_identifiers = [ identifier for identifier in identifiers if identifier_auto_merge_enabled(identifier, policy) ] matched_entity_ids: set[int] = set() for identifier in enabled_identifiers: matched_entity_id = active_entity_id_for_resolution_key(connection, identifier) if matched_entity_id is not None: matched_entity_ids.add(matched_entity_id) if len(matched_entity_ids) == 1: entity_id = next(iter(matched_entity_ids)) entity_row = connection.execute( "SELECT entity_type FROM entities WHERE id = ?", (entity_id,), ).fetchone() if entity_row is not None and entity_types_conflict(entity_row["entity_type"], candidate.get("entity_type")): entity_id = create_entity_for_candidate(connection, candidate) enabled_identifiers = [] else: entity_id = create_entity_for_candidate(connection, candidate) if len(matched_entity_ids) > 1: enabled_identifiers = [] inserted_identifier_ids: dict[int, dict[str, object]] = {} for identifier in identifiers: identifier_id = ensure_entity_identifier(connection, entity_id=entity_id, identifier=identifier) inserted_identifier_ids[identifier_id] = identifier for identifier_id, identifier in inserted_identifier_ids.items(): if identifier in enabled_identifiers: ensure_entity_resolution_key( connection, entity_id=entity_id, identifier_id=identifier_id, identifier=identifier, ) recompute_entity_caches(connection, entity_id) return entity_id def entity_display_label_from_row(row: sqlite3.Row | dict[str, object] | None) -> str: if row is None: return "Unknown Entity" display_name = normalize_entity_text(row["display_name"] if "display_name" in row.keys() else None) # type: ignore[attr-defined,index] primary_email = normalize_entity_email(row["primary_email"] if "primary_email" in row.keys() else None) # type: ignore[attr-defined,index] primary_phone = normalize_entity_text(row["primary_phone"] if "primary_phone" in row.keys() else None) # type: ignore[attr-defined,index] entity_id = row["id"] if "id" in row.keys() else None # type: ignore[attr-defined,index] if display_name and primary_email: return f"{display_name} <{primary_email}>" if display_name: return display_name if primary_email: return primary_email if primary_phone: return primary_phone return f"Unknown Entity {entity_id}" if entity_id is not None else "Unknown Entity" def entity_display_label(connection: sqlite3.Connection, entity_id: int) -> str: row = connection.execute( """ SELECT * FROM entities WHERE id = ? """, (entity_id,), ).fetchone() return entity_display_label_from_row(row) def rebuild_document_entity_caches(connection: sqlite3.Connection, document_id: int) -> dict[str, object]: document_row = connection.execute( f""" SELECT id, {MANUAL_FIELD_LOCKS_COLUMN} AS locks_json FROM documents WHERE id = ? """, (document_id,), ).fetchone() if document_row is None: raise RetrieverError(f"Unknown document id: {document_id}") locked_fields = set(normalize_string_list(document_row["locks_json"])) rows = connection.execute( """ SELECT de.role, de.ordinal, e.* FROM document_entities de JOIN entities e ON e.id = de.entity_id WHERE de.document_id = ? AND e.canonical_status = ? ORDER BY de.role ASC, de.ordinal ASC, de.id ASC """, (document_id, ENTITY_STATUS_ACTIVE), ).fetchall() labels_by_role: dict[str, list[str]] = defaultdict(list) for row in rows: label = entity_display_label_from_row(row) if label not in labels_by_role[row["role"]]: labels_by_role[row["role"]].append(label) updates: dict[str, object] = {} if "author" not in locked_fields: updates["author"] = labels_by_role.get("author", [None])[0] if labels_by_role.get("author") else None if "participants" not in locked_fields: updates["participants"] = ", ".join(labels_by_role.get("participant", [])) or None if "recipients" not in locked_fields: updates["recipients"] = ", ".join(labels_by_role.get("recipient", [])) or None if "custodian" not in locked_fields: updates["custodians_json"] = json.dumps(labels_by_role.get("custodian", []), ensure_ascii=True) if updates: set_clause = ", ".join(f"{quote_identifier(column)} = ?" for column in updates) connection.execute( f""" UPDATE documents SET {set_clause}, updated_at = ? WHERE id = ? """, [*updates.values(), utc_now(), document_id], ) return updates def entity_candidate_is_globally_ignored(connection: sqlite3.Connection, candidate: dict[str, object]) -> bool: row = connection.execute( """ SELECT 1 FROM entity_overrides WHERE scope_type = 'global' AND override_effect = 'ignore' AND ( normalized_candidate_key = ? OR source_entity_id IS NULL AND normalized_candidate_key IS NULL AND source_hint = ? ) LIMIT 1 """, (candidate.get("normalized_candidate_key"), candidate.get("raw_value")), ).fetchone() return row is not None def document_entity_override_for_candidate( connection: sqlite3.Connection, *, document_id: int, role: str, source_entity_id: int, candidate: dict[str, object], ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM entity_overrides WHERE scope_type = 'document' AND scope_id = ? AND (role IS NULL OR role = ?) AND ( source_entity_id = ? OR normalized_candidate_key = ? OR ( source_entity_id IS NULL AND normalized_candidate_key IS NULL AND source_hint = ? ) ) ORDER BY id DESC LIMIT 1 """, ( int(document_id), role, int(source_entity_id), candidate.get("normalized_candidate_key"), candidate.get("raw_value"), ), ).fetchone() def sync_document_entities( connection: sqlite3.Connection, document_id: int, *, refresh_fts: bool = True, ) -> dict[str, object]: if not all( table_exists(connection, table_name) for table_name in ("entities", "entity_identifiers", "entity_resolution_keys", "document_entities") ): return {"document_id": document_id, "synced": False, "reason": "entity schema unavailable"} active_rows = active_occurrence_rows_for_document(connection, document_id) auto_links: list[tuple[int, int, str, int, str, str | None, str, str, str]] = [] seen_role_entities: set[tuple[str, int]] = set() ordinals_by_role: dict[str, int] = defaultdict(int) for occurrence_row in occurrence_rows_in_preferred_order(active_rows): dataset_source_row = dataset_source_row_for_occurrence(connection, occurrence_row) if dataset_source_row is None: dataset_source_row = dataset_source_row_for_document_membership(connection, document_id) if dataset_source_row is None: continue policy = source_backed_dataset_policy_for_source(connection, dataset_source_row) role_values = ( ("author", occurrence_row["extracted_author"]), ("participant", occurrence_row["extracted_participants"]), ("recipient", occurrence_row["extracted_recipients"]), ("custodian", occurrence_row["custodian"]), ) raw_entity_hints = occurrence_row["entity_hints_json"] if "entity_hints_json" in occurrence_row.keys() else None for role, raw_value in role_values: for candidate in parse_entity_candidates_with_hints(raw_value, role=role, raw_hints=raw_entity_hints): if entity_candidate_is_globally_ignored(connection, candidate): continue entity_id = resolve_entity_candidate(connection, candidate, policy=policy) override_row = document_entity_override_for_candidate( connection, document_id=document_id, role=role, source_entity_id=entity_id, candidate=candidate, ) if override_row is not None: if override_row["override_effect"] == "remove": continue if override_row["override_effect"] == "replace": replacement_entity_id = override_row["replacement_entity_id"] if replacement_entity_id is None: continue canonical_replacement_id = canonicalize_entity_id(connection, int(replacement_entity_id)) if canonical_replacement_id is None: continue replacement_row = connection.execute( """ SELECT canonical_status FROM entities WHERE id = ? """, (canonical_replacement_id,), ).fetchone() if replacement_row is None or replacement_row["canonical_status"] != ENTITY_STATUS_ACTIVE: continue entity_id = canonical_replacement_id role_entity_key = (role, entity_id) if role_entity_key in seen_role_entities: continue seen_role_entities.add(role_entity_key) ordinal = ordinals_by_role[role] ordinals_by_role[role] += 1 evidence = { "raw_value": candidate.get("raw_value"), "occurrence_id": int(occurrence_row["id"]), "dataset_source_id": int(dataset_source_row["id"]) if dataset_source_row is not None else None, "normalized_candidate_key": candidate.get("normalized_candidate_key"), } auto_links.append( ( document_id, entity_id, role, ordinal, "auto", None, json.dumps(evidence, ensure_ascii=True, sort_keys=True), utc_now(), utc_now(), ) ) connection.execute( """ DELETE FROM document_entities WHERE document_id = ? AND assignment_mode = 'auto' """, (document_id,), ) if auto_links: connection.executemany( """ INSERT OR IGNORE INTO document_entities ( document_id, entity_id, role, ordinal, assignment_mode, observed_title, evidence_json, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, auto_links, ) cache_updates = rebuild_document_entity_caches(connection, document_id) if refresh_fts: refresh_documents_fts_row(connection, document_id) return { "document_id": document_id, "synced": True, "auto_link_count": len(auto_links), "cache_updates": cache_updates, } def refresh_document_control_number_aliases(connection: sqlite3.Connection, document_id: int) -> None: now = utc_now() connection.execute( "DELETE FROM document_control_number_aliases WHERE document_id = ?", (document_id,), ) document_row = connection.execute( "SELECT control_number FROM documents WHERE id = ?", (document_id,), ).fetchone() alias_rows: list[tuple[int, int | None, str, str, int, str, str]] = [] if document_row is not None and normalize_whitespace(str(document_row["control_number"] or "")): alias_rows.append( ( document_id, None, str(document_row["control_number"]), "document_primary", 1, now, now, ) ) occurrence_rows = connection.execute( """ SELECT id, occurrence_control_number, lifecycle_status FROM document_occurrences WHERE document_id = ? ORDER BY id ASC """, (document_id,), ).fetchall() for occurrence_row in occurrence_rows: alias_value = normalize_whitespace(str(occurrence_row["occurrence_control_number"] or "")) if not alias_value: continue alias_rows.append( ( document_id, int(occurrence_row["id"]), alias_value, "occurrence_control_number", 1 if occurrence_row["lifecycle_status"] == ACTIVE_OCCURRENCE_STATUS else 0, now, now, ) ) if alias_rows: connection.executemany( """ INSERT INTO document_control_number_aliases ( document_id, occurrence_id, alias_value, alias_type, active_flag, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, alias_rows, ) def refresh_canonical_metadata_conflicts( connection: sqlite3.Connection, document_id: int, active_rows: list[sqlite3.Row], ) -> None: tracked_fields = { "author": "extracted_author", "title": "extracted_title", "subject": "extracted_subject", "participants": "extracted_participants", "recipients": "extracted_recipients", "date_created": "extracted_doc_authored_at", "date_modified": "extracted_doc_modified_at", "content_type": "extracted_content_type", } now = utc_now() connection.execute( "DELETE FROM canonical_metadata_conflicts WHERE document_id = ?", (document_id,), ) rows_to_insert: list[tuple[int, str, int, str, str, str]] = [] for field_name, occurrence_column in tracked_fields.items(): distinct_values = { normalize_whitespace(str(row[occurrence_column] or "")): row for row in active_rows if normalize_whitespace(str(row[occurrence_column] or "")) } if len(distinct_values) <= 1: continue for value, row in distinct_values.items(): rows_to_insert.append( ( document_id, field_name, int(row["id"]), value, now, now, ) ) if rows_to_insert: connection.executemany( """ INSERT INTO canonical_metadata_conflicts ( document_id, field_name, occurrence_id, value, first_seen_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?) """, rows_to_insert, ) def dataset_source_row_for_occurrence_values( connection: sqlite3.Connection, *, source_kind: object, source_rel_path: object, production_id: object = None, ) -> sqlite3.Row | None: normalized_source_kind = normalize_whitespace(str(source_kind or "")).lower() normalized_source_rel_path = normalize_whitespace(str(source_rel_path or "")) if normalized_source_kind == FILESYSTEM_SOURCE_KIND: return get_dataset_source_row( connection, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), ) if normalized_source_kind in {PST_SOURCE_KIND, MBOX_SOURCE_KIND} and normalized_source_rel_path: return get_dataset_source_row( connection, source_kind=normalized_source_kind, source_locator=normalized_source_rel_path, ) if normalized_source_kind == PRODUCTION_SOURCE_KIND: if production_id is None: return None production_row = connection.execute( "SELECT rel_root FROM productions WHERE id = ?", (production_id,), ).fetchone() if production_row is None: return None return get_dataset_source_row( connection, source_kind=PRODUCTION_SOURCE_KIND, source_locator=str(production_row["rel_root"]), ) if normalized_source_kind == SLACK_EXPORT_SOURCE_KIND and normalized_source_rel_path: return connection.execute( """ SELECT * FROM dataset_sources WHERE source_kind = ? AND (? = source_locator OR ? LIKE source_locator || '/%') ORDER BY LENGTH(source_locator) DESC, id ASC LIMIT 1 """, (SLACK_EXPORT_SOURCE_KIND, normalized_source_rel_path, normalized_source_rel_path), ).fetchone() return None def dataset_source_row_for_occurrence( connection: sqlite3.Connection, occurrence_row: sqlite3.Row, ) -> sqlite3.Row | None: if "dataset_source_id" in occurrence_row.keys() and occurrence_row["dataset_source_id"] is not None: row = connection.execute( """ SELECT * FROM dataset_sources WHERE id = ? """, (int(occurrence_row["dataset_source_id"]),), ).fetchone() if row is not None: return row return dataset_source_row_for_occurrence_values( connection, source_kind=occurrence_row["source_kind"], source_rel_path=occurrence_row["source_rel_path"], production_id=occurrence_row["production_id"], ) def dataset_source_row_for_document_membership( connection: sqlite3.Connection, document_id: int, ) -> sqlite3.Row | None: return connection.execute( """ SELECT ds.* FROM dataset_documents dd JOIN dataset_sources ds ON ds.id = dd.dataset_source_id WHERE dd.document_id = ? AND dd.dataset_source_id IS NOT NULL ORDER BY dd.dataset_id ASC, dd.dataset_source_id ASC LIMIT 1 """, (int(document_id),), ).fetchone() def refresh_source_backed_dataset_memberships_for_document(connection: sqlite3.Connection, document_id: int) -> None: active_rows = active_occurrence_rows_for_document(connection, document_id) document_row = connection.execute( "SELECT parent_document_id FROM documents WHERE id = ?", (document_id,), ).fetchone() connection.execute( """ DELETE FROM dataset_documents WHERE document_id = ? AND dataset_source_id IS NOT NULL """, (document_id,), ) if document_row is not None and document_row["parent_document_id"] is not None: parent_source_rows = connection.execute( """ SELECT dataset_id, dataset_source_id FROM dataset_documents WHERE document_id = ? AND dataset_source_id IS NOT NULL ORDER BY dataset_id ASC, dataset_source_id ASC """, (document_row["parent_document_id"],), ).fetchall() for parent_source_row in parent_source_rows: ensure_dataset_document_membership( connection, dataset_id=int(parent_source_row["dataset_id"]), document_id=document_id, dataset_source_id=int(parent_source_row["dataset_source_id"]), ) refresh_document_dataset_cache(connection, document_id) return for occurrence_row in active_rows: dataset_source_row = dataset_source_row_for_occurrence(connection, occurrence_row) if dataset_source_row is None: continue ensure_dataset_document_membership( connection, dataset_id=int(dataset_source_row["dataset_id"]), document_id=document_id, dataset_source_id=int(dataset_source_row["id"]), ) refresh_document_dataset_cache(connection, document_id) def refresh_document_from_occurrences(connection: sqlite3.Connection, document_id: int) -> dict[str, object]: document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (document_id,), ).fetchone() if document_row is None: raise RetrieverError(f"Unknown document id: {document_id}") active_rows = active_occurrence_rows_for_document(connection, document_id) if not active_rows: if document_row["canonical_status"] == CANONICAL_STATUS_MERGED: connection.execute( """ UPDATE documents SET dataset_id = NULL, updated_at = ? WHERE id = ? """, (utc_now(), document_id), ) refresh_document_control_number_aliases(connection, document_id) return { "document_id": document_id, "preferred_occurrence_id": None, "active_occurrence_count": 0, "canonical_status": CANONICAL_STATUS_MERGED, } connection.execute( """ UPDATE documents SET canonical_status = ?, lifecycle_status = ?, dataset_id = NULL, updated_at = ? WHERE id = ? """, (CANONICAL_STATUS_DERELICT, "missing", utc_now(), document_id), ) refresh_document_control_number_aliases(connection, document_id) refresh_documents_fts_row(connection, document_id) return { "document_id": document_id, "preferred_occurrence_id": None, "active_occurrence_count": 0, "canonical_status": CANONICAL_STATUS_DERELICT, } preferred_row = select_preferred_occurrence(active_rows) assert preferred_row is not None locked_fields = set(normalize_string_list(document_row[MANUAL_FIELD_LOCKS_COLUMN])) resolved_author = occurrence_field_value(preferred_row, active_rows, "extracted_author") resolved_content_type = occurrence_field_value(preferred_row, active_rows, "extracted_content_type") resolved_custodians = custodian_values_from_occurrence_rows(active_rows) resolved_date_created = occurrence_field_value(preferred_row, active_rows, "extracted_doc_authored_at") resolved_date_modified = occurrence_field_value(preferred_row, active_rows, "extracted_doc_modified_at") resolved_title = occurrence_field_value(preferred_row, active_rows, "extracted_title") resolved_subject = occurrence_field_value(preferred_row, active_rows, "extracted_subject") resolved_participants = occurrence_field_value(preferred_row, active_rows, "extracted_participants") resolved_recipients = occurrence_field_value(preferred_row, active_rows, "extracted_recipients") if "author" in locked_fields: resolved_author = document_row["author"] if "content_type" in locked_fields: resolved_content_type = document_row["content_type"] if "date_created" in locked_fields: resolved_date_created = document_row["date_created"] if "date_modified" in locked_fields: resolved_date_modified = document_row["date_modified"] if "title" in locked_fields: resolved_title = document_row["title"] if "subject" in locked_fields: resolved_subject = document_row["subject"] if "participants" in locked_fields: resolved_participants = document_row["participants"] if "recipients" in locked_fields: resolved_recipients = document_row["recipients"] canonical_control_number = None for row in sorted( active_rows, key=lambda candidate: ( 0 if int(candidate["id"]) == int(preferred_row["id"]) else 1, source_kind_priority(candidate["source_kind"]), text_status_priority(candidate["text_status"]), parse_utc_timestamp(candidate["ingested_at"]) or datetime.max.replace(tzinfo=timezone.utc), int(candidate["id"]), ), ): candidate_control_number = normalize_whitespace(str(row["occurrence_control_number"] or "")) if candidate_control_number: canonical_control_number = candidate_control_number break updated_values = { "control_number": canonical_control_number or document_row["control_number"], "canonical_kind": canonical_kind_from_metadata( extracted_content_type=occurrence_field_value(preferred_row, active_rows, "extracted_content_type"), extracted_kind=occurrence_field_value(preferred_row, active_rows, "extracted_kind"), file_type=preferred_row["file_type"], source_kind=preferred_row["source_kind"], ), "canonical_status": CANONICAL_STATUS_ACTIVE, "merged_into_document_id": None, "source_kind": preferred_row["source_kind"], "source_rel_path": preferred_row["source_rel_path"], "source_item_id": preferred_row["source_item_id"], "source_folder_path": preferred_row["source_folder_path"], "production_id": preferred_row["production_id"], "begin_bates": preferred_row["begin_bates"], "end_bates": preferred_row["end_bates"], "begin_attachment": preferred_row["begin_attachment"], "end_attachment": preferred_row["end_attachment"], "rel_path": preferred_row["rel_path"], "file_name": preferred_row["file_name"], "file_type": preferred_row["file_type"], "file_size": preferred_row["file_size"], "author": resolved_author, "content_type": resolved_content_type, "custodians_json": json.dumps(resolved_custodians), "date_created": resolved_date_created, "date_modified": resolved_date_modified, "title": resolved_title, "subject": resolved_subject, "participants": resolved_participants, "recipients": resolved_recipients, "file_hash": preferred_row["file_hash"], "text_status": min((row["text_status"] for row in active_rows), key=text_status_priority), "lifecycle_status": "active", "ingested_at": min( (parse_utc_timestamp(row["ingested_at"]) or datetime.max.replace(tzinfo=timezone.utc) for row in active_rows) ).isoformat().replace("+00:00", "Z"), "last_seen_at": max( ( parse_utc_timestamp(row["last_seen_at"]) or datetime.min.replace(tzinfo=timezone.utc) for row in active_rows ) ).isoformat().replace("+00:00", "Z"), "updated_at": utc_now(), } connection.execute( """ UPDATE documents SET control_number = ?, canonical_kind = ?, canonical_status = ?, merged_into_document_id = ?, source_kind = ?, source_rel_path = ?, source_item_id = ?, source_folder_path = ?, production_id = ?, begin_bates = ?, end_bates = ?, begin_attachment = ?, end_attachment = ?, rel_path = ?, file_name = ?, file_type = ?, file_size = ?, author = ?, content_type = ?, custodians_json = ?, date_created = ?, date_modified = ?, title = ?, subject = ?, participants = ?, recipients = ?, file_hash = ?, text_status = ?, lifecycle_status = ?, ingested_at = ?, last_seen_at = ?, updated_at = ? WHERE id = ? """, ( updated_values["control_number"], updated_values["canonical_kind"], updated_values["canonical_status"], updated_values["merged_into_document_id"], updated_values["source_kind"], updated_values["source_rel_path"], updated_values["source_item_id"], updated_values["source_folder_path"], updated_values["production_id"], updated_values["begin_bates"], updated_values["end_bates"], updated_values["begin_attachment"], updated_values["end_attachment"], updated_values["rel_path"], updated_values["file_name"], updated_values["file_type"], updated_values["file_size"], updated_values["author"], updated_values["content_type"], updated_values["custodians_json"], updated_values["date_created"], updated_values["date_modified"], updated_values["title"], updated_values["subject"], updated_values["participants"], updated_values["recipients"], updated_values["file_hash"], updated_values["text_status"], updated_values["lifecycle_status"], updated_values["ingested_at"], updated_values["last_seen_at"], updated_values["updated_at"], document_id, ), ) refresh_canonical_metadata_conflicts(connection, document_id, active_rows) refresh_document_control_number_aliases(connection, document_id) sync_document_entities(connection, document_id, refresh_fts=False) refresh_documents_fts_row(connection, document_id) return { "document_id": document_id, "preferred_occurrence_id": int(preferred_row["id"]), "active_occurrence_count": len(active_rows), "canonical_status": CANONICAL_STATUS_ACTIVE, } def get_document_by_dedupe_key( connection: sqlite3.Connection, *, basis: str, key_value: str | None, ) -> sqlite3.Row | None: normalized_key = normalize_whitespace(str(key_value or "")) if not normalized_key: return None return connection.execute( """ SELECT d.* FROM document_dedupe_keys dk JOIN documents d ON d.id = dk.document_id WHERE dk.basis = ? AND dk.key_value = ? """, (basis, normalized_key), ).fetchone() def bind_document_dedupe_key( connection: sqlite3.Connection, *, basis: str, key_value: str | None, document_id: int, ) -> bool: normalized_key = normalize_whitespace(str(key_value or "")) if not normalized_key: return False now = utc_now() cursor = connection.execute( """ INSERT INTO document_dedupe_keys (basis, key_value, document_id, created_at, updated_at) VALUES (?, ?, ?, ?, ?) ON CONFLICT(basis, key_value) DO UPDATE SET document_id = excluded.document_id, updated_at = excluded.updated_at WHERE document_dedupe_keys.document_id != excluded.document_id """, (basis, normalized_key, document_id, now, now), ) return int(cursor.rowcount or 0) > 0 def find_active_occurrence_by_source_identity( connection: sqlite3.Connection, *, source_kind: str | None, custodian: str | None, source_rel_path: str | None, source_item_id: str | None, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM document_occurrences WHERE source_kind = ? AND COALESCE(custodian, '') = COALESCE(?, '') AND COALESCE(source_rel_path, '') = COALESCE(?, '') AND COALESCE(source_item_id, '') = COALESCE(?, '') AND lifecycle_status = 'active' ORDER BY id ASC LIMIT 1 """, (source_kind, custodian, source_rel_path, source_item_id), ).fetchone() def container_root_occurrence_rows_for_source( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, include_deleted: bool = False, ) -> list[sqlite3.Row]: normalized_source_kind = normalize_whitespace(str(source_kind or "")).lower() normalized_source_rel_path = normalize_whitespace(str(source_rel_path or "")) if not normalized_source_kind or not normalized_source_rel_path: return [] clauses = [ "parent_occurrence_id IS NULL", "source_kind = ?", "source_rel_path = ?", ] parameters: list[object] = [normalized_source_kind, normalized_source_rel_path] if not include_deleted: clauses.append("lifecycle_status != 'deleted'") return connection.execute( f""" SELECT * FROM document_occurrences WHERE {' AND '.join(clauses)} ORDER BY id ASC """, parameters, ).fetchall() def container_document_ids_for_root_occurrence_ids( connection: sqlite3.Connection, root_occurrence_ids: list[int], ) -> set[int]: normalized_ids = sorted({int(occurrence_id) for occurrence_id in root_occurrence_ids}) if not normalized_ids: return set() placeholders = ", ".join("?" for _ in normalized_ids) rows = connection.execute( f""" SELECT DISTINCT document_id FROM document_occurrences WHERE id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders}) ORDER BY document_id ASC """, [*normalized_ids, *normalized_ids], ).fetchall() return {int(row["document_id"]) for row in rows} def container_root_document_ids_for_source( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, include_deleted: bool = False, ) -> set[int]: return { int(row["document_id"]) for row in container_root_occurrence_rows_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, include_deleted=include_deleted, ) } def container_document_ids_for_source( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, include_deleted: bool = False, ) -> set[int]: root_occurrence_ids = [ int(row["id"]) for row in container_root_occurrence_rows_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, include_deleted=include_deleted, ) ] return container_document_ids_for_root_occurrence_ids(connection, root_occurrence_ids) def upsert_document_occurrence( connection: sqlite3.Connection, *, document_id: int, existing_occurrence_id: int | None, parent_occurrence_id: int | None, occurrence_control_number: str | None, source_kind: str | None, source_rel_path: str | None, source_item_id: str | None, source_folder_path: str | None, production_id: int | None, begin_bates: str | None, end_bates: str | None, begin_attachment: str | None, end_attachment: str | None, rel_path: str, file_name: str, file_type: str | None, mime_type: str | None, file_size: int | None, file_hash: str | None, custodian: str | None, fs_created_at: str | None, fs_modified_at: str | None, extracted: dict[str, object], has_preview: bool, text_status: str, ingested_at: str, last_seen_at: str, updated_at: str, ) -> int: dataset_source_row = dataset_source_row_for_occurrence_values( connection, source_kind=source_kind, source_rel_path=source_rel_path, production_id=production_id, ) dataset_source_id = int(dataset_source_row["id"]) if dataset_source_row is not None else None raw_entity_hints = extracted.get("entity_hints") entity_hints_json = json.dumps( raw_entity_hints if isinstance(raw_entity_hints, dict) else {}, ensure_ascii=True, sort_keys=True, ) occurrence_values = ( document_id, dataset_source_id, parent_occurrence_id, occurrence_control_number, source_kind, source_rel_path, source_item_id, source_folder_path, production_id, begin_bates, end_bates, begin_attachment, end_attachment, rel_path, file_name, file_type, mime_type, file_size, file_hash, custodian, fs_created_at, fs_modified_at, extracted.get("author"), extracted.get("title"), extracted.get("subject"), extracted.get("participants"), extracted.get("recipients"), extracted.get("date_created"), extracted.get("date_modified"), extracted.get("content_type"), canonical_kind_from_metadata( extracted_content_type=extracted.get("content_type"), file_type=file_type, source_kind=source_kind, ), entity_hints_json, text_status, ACTIVE_OCCURRENCE_STATUS, 1 if has_preview else 0, ingested_at, last_seen_at, updated_at, ) if existing_occurrence_id is None: connection.execute( """ INSERT INTO document_occurrences ( document_id, dataset_source_id, parent_occurrence_id, occurrence_control_number, source_kind, source_rel_path, source_item_id, source_folder_path, production_id, begin_bates, end_bates, begin_attachment, end_attachment, rel_path, file_name, file_type, mime_type, file_size, file_hash, custodian, fs_created_at, fs_modified_at, extracted_author, extracted_title, extracted_subject, extracted_participants, extracted_recipients, extracted_doc_authored_at, extracted_doc_modified_at, extracted_content_type, extracted_kind, entity_hints_json, text_status, lifecycle_status, has_preview, ingested_at, last_seen_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, occurrence_values, ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) connection.execute( """ UPDATE document_occurrences SET document_id = ?, dataset_source_id = ?, parent_occurrence_id = ?, occurrence_control_number = ?, source_kind = ?, source_rel_path = ?, source_item_id = ?, source_folder_path = ?, production_id = ?, begin_bates = ?, end_bates = ?, begin_attachment = ?, end_attachment = ?, rel_path = ?, file_name = ?, file_type = ?, mime_type = ?, file_size = ?, file_hash = ?, custodian = ?, fs_created_at = ?, fs_modified_at = ?, extracted_author = ?, extracted_title = ?, extracted_subject = ?, extracted_participants = ?, extracted_recipients = ?, extracted_doc_authored_at = ?, extracted_doc_modified_at = ?, extracted_content_type = ?, extracted_kind = ?, entity_hints_json = ?, text_status = ?, lifecycle_status = ?, has_preview = ?, ingested_at = ?, last_seen_at = ?, updated_at = ? WHERE id = ? """, (*occurrence_values, existing_occurrence_id), ) return existing_occurrence_id def infer_registry_field_type(sqlite_type: str | None) -> str: type_name = (sqlite_type or "").upper() if "DATE" in type_name or "TIME" in type_name: return "date" if "INT" in type_name: return "integer" if any(marker in type_name for marker in ("REAL", "FLOA", "DOUB")): return "real" return "text" def sanitize_field_name(field_name: str) -> str: sanitized = re.sub(r"[^a-zA-Z0-9_]+", "_", field_name.strip()).strip("_").lower() if not sanitized: raise RetrieverError("Field name becomes empty after sanitization.") if sanitized[0].isdigit(): sanitized = f"field_{sanitized}" if sanitized in BUILTIN_FIELD_TYPES: raise RetrieverError(f"Field name '{sanitized}' conflicts with a built-in document column.") if sanitized in INTERNAL_DOCUMENT_COLUMNS: raise RetrieverError(f"Field name '{sanitized}' conflicts with a system-managed document column.") return sanitized def parse_pdf_date(value: object) -> str | None: if not value or not isinstance(value, str): return None raw = value.strip() if not raw: return None if raw.startswith("D:"): raw = raw[2:] if not re.fullmatch(r"\d{4}(?:\d{2}){0,5}(?:Z|[+\-]\d{2}'?\d{2}'?)?", raw): return None match = re.match( r"^(?P<year>\d{4})(?P<month>\d{2})?(?P<day>\d{2})?(?P<hour>\d{2})?(?P<minute>\d{2})?(?P<second>\d{2})?", raw, ) if not match: return None parts = match.groupdict(default=None) month = int(parts["month"] or "1") day = int(parts["day"] or "1") hour = int(parts["hour"] or "0") minute = int(parts["minute"] or "0") second = int(parts["second"] or "0") try: dt = datetime(int(parts["year"]), month, day, hour, minute, second, tzinfo=timezone.utc) except ValueError: return None return dt.isoformat().replace("+00:00", "Z") def parse_iso_datetime(value: str) -> str | None: raw = value.strip() if not raw: return None candidate = raw[:-1] + "+00:00" if raw.endswith("Z") else raw try: dt = datetime.fromisoformat(candidate) except ValueError: try: parsed_date = date.fromisoformat(raw) except ValueError: return None dt = datetime(parsed_date.year, parsed_date.month, parsed_date.day, tzinfo=timezone.utc) if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) else: dt = dt.astimezone(timezone.utc) return dt.replace(microsecond=0).isoformat().replace("+00:00", "Z") def normalize_date_field_value(value: object) -> str | None: if value in (None, ""): return None if isinstance(value, datetime): normalized = value if normalized.tzinfo is None: normalized = normalized.replace(tzinfo=timezone.utc) else: normalized = normalized.astimezone(timezone.utc) return normalized.replace(microsecond=0).isoformat().replace("+00:00", "Z") if isinstance(value, date): return value.isoformat() if not isinstance(value, str): return None raw = value.strip() if not raw: return None if re.fullmatch(r"\d{4}-\d{2}-\d{2}", raw): try: return date.fromisoformat(raw).isoformat() except ValueError: return None return parse_iso_datetime(raw) def normalize_datetime(value: object) -> str | None: if value in (None, ""): return None if isinstance(value, datetime): dt = value if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) else: dt = dt.astimezone(timezone.utc) return dt.replace(microsecond=0).isoformat().replace("+00:00", "Z") if isinstance(value, str): parsed = parse_iso_datetime(value) if parsed is not None: return parsed parsed = parse_pdf_date(value) if parsed is not None: return parsed for fmt in ( "%m/%d/%Y %I:%M:%S %p", "%m/%d/%Y %I:%M %p", "%m/%d/%y %I:%M:%S %p", "%m/%d/%y %I:%M %p", ): try: dt = datetime.strptime(value.strip(), fmt).replace(tzinfo=timezone.utc) return dt.isoformat().replace("+00:00", "Z") except ValueError: pass try: dt = parsedate_to_datetime(value) except (TypeError, ValueError, IndexError): return value.strip() or None if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) else: dt = dt.astimezone(timezone.utc) return dt.replace(microsecond=0).isoformat().replace("+00:00", "Z") return None def decode_bytes(data: bytes, declared_encoding: str | None = None) -> tuple[str, str, str | None]: def normalized_decoded_text(value: str) -> str: return value.lstrip("\ufeff") if declared_encoding: try: return normalized_decoded_text(data.decode(declared_encoding)), "ok", declared_encoding except Exception: pass try: return normalized_decoded_text(data.decode("utf-8-sig")), "ok", "utf-8" except UnicodeDecodeError: pass charset_normalizer_module = load_dependency("charset_normalizer") if charset_normalizer_module is not None: best = charset_normalizer_module.from_bytes(data).best() if best is not None: text = normalized_decoded_text(str(best)) status = "partial" if "\ufffd" in text else "ok" return text, status, best.encoding text = normalized_decoded_text(data.decode("utf-8", errors="replace")) status = "partial" if "\ufffd" in text else "ok" return text, status, None def strip_html_tags(text: str) -> str: without_scripts = re.sub(r"(?is)<(script|style).*?>.*?</\1>", " ", text) with_breaks = re.sub(r"(?is)<br\s*/?>", "\n", without_scripts) with_breaks = re.sub( r"(?is)</(?:p|div|li|tr|td|th|h[1-6]|section|article|blockquote|pre|ul|ol|table)>", "\n", with_breaks, ) without_tags = re.sub(r"(?s)<[^>]+>", " ", with_breaks) return normalize_whitespace(html.unescape(without_tags)) def normalize_participant_token(value: str | None) -> str | None: if value is None: return None normalized = normalize_whitespace(value) normalized = re.sub(r"\s*<\s*", " <", normalized) normalized = re.sub(r"\s*>\s*", ">", normalized) normalized = normalized.strip(" ,;") return normalized or None def append_unique_participants( participants: list[str], seen: set[str], raw_values: list[str | None], ) -> None: for raw_value in raw_values: if not raw_value: continue normalized_candidate_text = normalize_participant_token(raw_value) if not normalized_candidate_text: continue if "@" not in normalized_candidate_text: for raw_part in re.split(r"\s*;\s*|\n+", normalized_candidate_text): rendered = normalize_participant_token(raw_part) if not rendered: continue key = rendered.lower() if key not in seen: seen.add(key) participants.append(rendered) continue parsed_values = getaddresses([normalized_candidate_text.replace(";", ",")]) for display_name, email_address in parsed_values: normalized_name = normalize_participant_token(display_name) normalized_email = normalize_participant_token(email_address.lower() if email_address else None) if normalized_email and "@" in normalized_email: rendered = f"{normalized_name} <{normalized_email}>" if normalized_name and normalized_name.lower() != normalized_email else normalized_email elif normalized_name and not normalized_email: rendered = normalized_name else: rendered = None if not rendered: continue key = rendered.lower() if key not in seen: seen.add(key) participants.append(rendered) def sorted_unique_display_names(raw_values: list[object]) -> list[str]: unique_names: dict[str, str] = {} for raw_value in raw_values: normalized = normalize_participant_token(raw_value) if not normalized: continue unique_names.setdefault(normalized.casefold(), normalized) return sorted(unique_names.values(), key=str.casefold) def render_display_name_list( raw_values: list[object], *, max_names: int | None = None, ) -> str | None: names = sorted_unique_display_names(raw_values) if not names: return None if max_names is not None and max_names > 0 and len(names) > max_names: remaining = len(names) - max_names return ", ".join(names[:max_names]) + f" +{remaining} more" return ", ".join(names) def render_display_name_title( raw_values: list[object], *, max_names: int | None = None, ) -> str | None: names = sorted_unique_display_names(raw_values) if not names: return None if max_names is not None and max_names > 0 and len(names) > max_names: remaining = len(names) - max_names return " / ".join(names[:max_names]) + f" +{remaining} more" return " / ".join(names) def email_headers_to_metadata(headers: dict[str, str]) -> dict[str, str | None]: recipients = ", ".join(headers[key] for key in ("to", "cc", "bcc") if headers.get(key)) or None subject = normalize_generated_document_title(headers.get("subject")) participants: list[str] = [] seen: set[str] = set() append_unique_participants( participants, seen, [headers.get("from"), headers.get("to"), headers.get("cc"), headers.get("bcc")], ) return { "author": headers.get("from") or None, "recipients": recipients, "participants": ", ".join(participants) or None, "date_created": normalize_datetime(headers.get("sent") or headers.get("date")), "subject": subject, "title": subject, } def attachment_list_looks_like_filenames(raw_value: str) -> bool: candidates = [ normalize_whitespace(part).strip(" \t\r\n'\"()[]{}") for part in re.split(r"\s*[;,]\s*", raw_value) if normalize_whitespace(part) ] if not candidates: candidate = normalize_whitespace(raw_value).strip(" \t\r\n'\"()[]{}") if not candidate: return False candidates = [candidate] filename_like = 0 for candidate in candidates: leaf = candidate.replace("\\", "/").rsplit("/", 1)[-1] if re.search(r"\.[A-Za-z0-9]{1,8}$", leaf): filename_like += 1 return filename_like > 0 def normalize_generated_document_title(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")) or None if not normalized: return None match = ATTACHMENT_SUFFIX_PATTERN.match(normalized) if match is None or not attachment_list_looks_like_filenames(match.group("attachments")): return normalized trimmed_title = normalize_whitespace(match.group("title").rstrip(" -:;,")) return trimmed_title or normalized EMAIL_MESSAGE_ID_PATTERN = re.compile(r"<\s*([^<>]+?)\s*>|([^\s<>;,]+@[^\s<>;,]+)") EMAIL_THREAD_PREFIX_PATTERN = re.compile(r"^(?:(?:re|fw|fwd)\s*(?:\[\d+\])?\s*:\s*)+", flags=re.IGNORECASE) def extract_email_message_ids(value: object) -> list[str]: normalized = normalize_whitespace(str(value or "")) if not normalized: return [] message_ids: list[str] = [] seen: set[str] = set() for match in EMAIL_MESSAGE_ID_PATTERN.finditer(normalized): raw = normalize_whitespace(match.group(1) or match.group(2) or "") if not raw: continue normalized_id = raw.strip("<>").strip().lower() if not normalized_id or normalized_id in seen: continue seen.add(normalized_id) message_ids.append(normalized_id) return message_ids def normalize_email_message_id(value: object) -> str | None: message_ids = extract_email_message_ids(value) return message_ids[0] if message_ids else None def normalize_email_thread_subject(value: object, *, preserve_case: bool = False) -> str | None: normalized = normalize_whitespace(str(value or "")) if not normalized: return None previous = None current = normalized while previous != current: previous = current current = EMAIL_THREAD_PREFIX_PATTERN.sub("", current).strip() current = normalize_whitespace(current) if not current: return None return current if preserve_case else current.lower() def normalize_email_conversation_index_root(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")) if not normalized: return None compact = re.sub(r"\s+", "", normalized) if not compact: return None return compact[:44] if len(compact) > 44 else compact def email_participant_keys(author: object, recipients: object) -> set[str]: participants: list[str] = [] seen: set[str] = set() append_unique_participants( participants, seen, [ normalize_whitespace(str(author or "")) or None, normalize_whitespace(str(recipients or "")) or None, ], ) return {participant.lower() for participant in participants} def email_heuristic_scope_key(source_kind: object, source_rel_path: object) -> str: normalized_source_kind = normalize_whitespace(str(source_kind or "")).lower() normalized_source_rel_path = normalize_whitespace(str(source_rel_path or "")) if normalized_source_kind == MBOX_SOURCE_KIND and normalized_source_rel_path: return f"{MBOX_SOURCE_KIND}:{normalized_source_rel_path}" if normalized_source_kind == PST_SOURCE_KIND and normalized_source_rel_path: return f"{PST_SOURCE_KIND}:{normalized_source_rel_path}" if normalized_source_kind == FILESYSTEM_SOURCE_KIND: return f"{FILESYSTEM_SOURCE_KIND}:{filesystem_dataset_locator()}" return f"{normalized_source_kind or FILESYSTEM_SOURCE_KIND}:{normalized_source_rel_path or filesystem_dataset_locator()}" def extract_email_header_blocks(text: str, max_lines: int | None = None) -> list[dict[str, str]]: if not text: return [] lines = text.replace("\r\n", "\n").replace("\r", "\n").split("\n") recognized_keys = {"from", "to", "cc", "bcc", "sent", "date", "subject"} blocks: list[dict[str, str]] = [] headers: dict[str, str] = {} current_key: str | None = None started = False def flush_headers() -> None: nonlocal headers, current_key, started normalized_headers = { key: normalize_whitespace(value) for key, value in headers.items() if normalize_whitespace(value) } if "from" in normalized_headers and any( key in normalized_headers for key in ("to", "cc", "bcc", "subject", "sent", "date") ): blocks.append(dict(normalized_headers)) headers = {} current_key = None started = False for raw_line in lines[: max_lines or len(lines)]: stripped = raw_line.strip() if not stripped: if started and len(headers) >= 2: flush_headers() continue match = re.match(r"^(From|To|Cc|Bcc|Sent|Date|Subject):\s*(.*)$", stripped, flags=re.IGNORECASE) if match: key = match.group(1).lower() if started and key == "from" and len(headers) >= 2: flush_headers() current_key = key headers[current_key] = normalize_whitespace(match.group(2)) started = True continue if started and current_key and raw_line != raw_line.lstrip(): headers[current_key] = normalize_whitespace(f"{headers.get(current_key, '')} {stripped}") continue if started and len(headers) >= 2: flush_headers() if started: flush_headers() return [block for block in blocks if set(block).issubset(recognized_keys)] def extract_email_like_headers(text: str) -> dict[str, str | None]: blocks = extract_email_header_blocks(text, max_lines=60) if not blocks: return {} return email_headers_to_metadata(blocks[0]) def extract_email_chain_participants( text: str, initial_values: list[str | None] | None = None, ) -> str | None: participants: list[str] = [] seen: set[str] = set() append_unique_participants(participants, seen, list(initial_values or [])) for headers in extract_email_header_blocks(text): append_unique_participants( participants, seen, [headers.get("from"), headers.get("to"), headers.get("cc"), headers.get("bcc")], ) return ", ".join(participants) or None ICALENDAR_FILE_TYPES = {"ics", "ifb", "vcal", "vcs"} ICALENDAR_URL_PATTERN = re.compile(r"https?://[^\s<>'\"]+") CALENDAR_INVITE_TEXT_BLOCK_START = "[[RETRIEVER_CALENDAR_INVITE]]" CALENDAR_INVITE_TEXT_BLOCK_END = "[[/RETRIEVER_CALENDAR_INVITE]]" CALENDAR_INVITE_TEXT_BLOCK_PATTERN = re.compile( re.escape(CALENDAR_INVITE_TEXT_BLOCK_START) + r"\n(.*?)" + re.escape(CALENDAR_INVITE_TEXT_BLOCK_END) + r"\s*", re.DOTALL, ) def unfold_icalendar_lines(text: str) -> list[str]: unfolded: list[str] = [] for raw_line in str(text or "").replace("\r\n", "\n").replace("\r", "\n").split("\n"): if raw_line.startswith((" ", "\t")) and unfolded: unfolded[-1] += raw_line[1:] else: unfolded.append(raw_line) return unfolded def split_icalendar_property_line(line: str) -> tuple[str, str] | None: in_quotes = False escaped = False for index, char in enumerate(line): if char == '"' and not escaped: in_quotes = not in_quotes elif char == ":" and not in_quotes: return line[:index], line[index + 1 :] if escaped: escaped = False elif char == "\\": escaped = True return None def split_icalendar_parameter_values(raw_value: str) -> list[str]: values: list[str] = [] current: list[str] = [] in_quotes = False escaped = False for char in raw_value: if char == '"' and not escaped: in_quotes = not in_quotes current.append(char) elif char == "," and not in_quotes: values.append("".join(current)) current = [] else: current.append(char) if escaped: escaped = False elif char == "\\": escaped = True values.append("".join(current)) return values def unescape_icalendar_text(value: object) -> str | None: normalized = str(value or "") if not normalized: return None return ( normalized .replace("\\N", "\n") .replace("\\n", "\n") .replace("\\,", ",") .replace("\\;", ";") .replace("\\\\", "\\") ) or None def parse_icalendar_parameters(raw_segments: list[str]) -> dict[str, list[str]]: parameters: dict[str, list[str]] = {} for raw_segment in raw_segments: if "=" not in raw_segment: continue raw_key, raw_value = raw_segment.split("=", 1) key = normalize_whitespace(raw_key).upper() if not key: continue values = [ normalize_whitespace(str(unescape_icalendar_text(part.strip().strip('"')) or "")) for part in split_icalendar_parameter_values(raw_value) ] values = [value for value in values if value] if values: parameters[key] = values return parameters def parse_icalendar_property_line(line: str) -> dict[str, object] | None: split_line = split_icalendar_property_line(line) if split_line is None: return None raw_head, raw_value = split_line head_parts = raw_head.split(";") name = normalize_whitespace(head_parts[0]).upper() if not name: return None return { "name": name, "params": parse_icalendar_parameters(head_parts[1:]), "value": raw_value, } def parse_icalendar_datetime_value( value: object, *, tzid: object = None, value_type: object = None, ) -> dict[str, object]: raw = normalize_whitespace(str(value or "")) if not raw: return {} normalized_value_type = normalize_whitespace(str(value_type or "")).upper() or None normalized_tzid = normalize_whitespace(str(tzid or "")) or None if normalized_value_type == "DATE" or re.fullmatch(r"\d{8}", raw): try: parsed_date = date(int(raw[:4]), int(raw[4:6]), int(raw[6:8])) except ValueError: return {"raw": raw} return { "raw": raw, "date": parsed_date, "all_day": True, "iso": parsed_date.isoformat(), "tz_label": normalized_tzid, } candidate = raw[:-1] if raw.endswith("Z") else raw parsed_dt = None for fmt in ("%Y%m%dT%H%M%S", "%Y%m%dT%H%M"): try: parsed_dt = datetime.strptime(candidate, fmt) break except ValueError: continue if parsed_dt is None: return {"raw": raw} tz_label = normalized_tzid if raw.endswith("Z"): parsed_dt = parsed_dt.replace(tzinfo=timezone.utc) tz_label = parsed_dt.tzname() or "UTC" elif normalized_tzid: try: parsed_dt = parsed_dt.replace(tzinfo=ZoneInfo(normalized_tzid)) tz_label = parsed_dt.tzname() or normalized_tzid except Exception: tz_label = normalized_tzid return { "raw": raw, "datetime": parsed_dt, "all_day": False, "iso": ( normalize_datetime(parsed_dt) if isinstance(parsed_dt, datetime) and parsed_dt.tzinfo is not None else parsed_dt.replace(microsecond=0).isoformat() ), "tz_label": tz_label, "tzid": normalized_tzid, } def format_calendar_preview_date(value: date) -> str: return value.strftime("%b %d, %Y").replace(" 0", " ") def format_calendar_preview_datetime(value: datetime) -> str: return value.strftime("%b %d, %Y %I:%M %p").replace(" 0", " ") def format_calendar_preview_time(value: datetime) -> str: return value.strftime("%I:%M %p").lstrip("0") def format_icalendar_event_range( start_info: dict[str, object] | None, end_info: dict[str, object] | None = None, ) -> str | None: if not start_info: return None if start_info.get("all_day"): start_date = start_info.get("date") if not isinstance(start_date, date): return normalize_whitespace(str(start_info.get("raw") or "")) or None end_date = end_info.get("date") if isinstance(end_info, dict) else None if isinstance(end_date, date) and end_date > start_date: inclusive_end = end_date - timedelta(days=1) if inclusive_end != start_date: return ( f"{format_calendar_preview_date(start_date)} - " f"{format_calendar_preview_date(inclusive_end)} (all day)" ) return f"{format_calendar_preview_date(start_date)} (all day)" start_dt = start_info.get("datetime") if not isinstance(start_dt, datetime): return normalize_whitespace(str(start_info.get("raw") or "")) or None end_dt = end_info.get("datetime") if isinstance(end_info, dict) else None if isinstance(end_dt, datetime) and start_dt.tzinfo is not None and end_dt.tzinfo is not None: end_dt = end_dt.astimezone(start_dt.tzinfo) label = format_calendar_preview_datetime(start_dt) if isinstance(end_dt, datetime): if start_dt.date() == end_dt.date(): label = f"{label} - {format_calendar_preview_time(end_dt)}" else: label = f"{label} - {format_calendar_preview_datetime(end_dt)}" tz_label = normalize_whitespace(str(start_info.get("tz_label") or "")) or ( start_dt.tzname() if start_dt.tzinfo is not None else None ) if tz_label: label = f"{label} {tz_label}" return label def humanize_icalendar_enum(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")).upper() if not normalized: return None mapping = { "ACCEPTED": "Accepted", "CANCEL": "Canceled", "CANCELLED": "Canceled", "CANCELED": "Canceled", "CONFIRMED": "Confirmed", "COUNTER": "Counter", "DECLINED": "Declined", "DECLINECOUNTER": "Declined Counter", "NEEDS-ACTION": "Needs Action", "PUBLISH": "Published", "REQUEST": "Request", "TENTATIVE": "Tentative", } return mapping.get(normalized) or normalized.replace("-", " ").replace("_", " ").title() def format_icalendar_participant(value: object, params: dict[str, list[str]] | None = None) -> str | None: normalized_value = normalize_participant_token(unescape_icalendar_text(value)) if normalized_value and normalized_value.lower().startswith("mailto:"): normalized_value = normalize_participant_token(normalized_value[7:]) if normalized_value and "@" in normalized_value: normalized_value = normalized_value.lower() cn_values = params.get("CN") if isinstance(params, dict) else None common_name = normalize_participant_token(unescape_icalendar_text(cn_values[0])) if cn_values else None if common_name and normalized_value and common_name.lower() != normalized_value.lower(): return f"{common_name} <{normalized_value}>" return common_name or normalized_value def first_icalendar_url(text: object) -> str | None: normalized = str(text or "") if not normalized: return None candidates = [ match.group(0).rstrip(").,;") for match in ICALENDAR_URL_PATTERN.finditer(normalized) ] if not candidates: return None preferred_domains = ("meet.google.com", "zoom.us", "teams.microsoft.com", "webex.com") for candidate in candidates: if any(domain in candidate for domain in preferred_domains): return candidate return candidates[0] def summarize_icalendar_invite_status(metadata: dict[str, object] | None) -> str | None: if not isinstance(metadata, dict): return None parts = [ humanize_icalendar_enum(metadata.get("method")), humanize_icalendar_enum(metadata.get("status")), ] parts = [part for part in parts if part] return " · ".join(parts) if parts else None def parse_icalendar_event_metadata(text: object) -> dict[str, object] | None: normalized_text = str(text or "") if not normalized_text: return None upper_text = normalized_text.upper() if "BEGIN:VCALENDAR" not in upper_text and "BEGIN:VEVENT" not in upper_text: return None calendar_properties: dict[str, list[dict[str, object]]] = defaultdict(list) event_properties: dict[str, list[dict[str, object]]] | None = None for raw_line in unfold_icalendar_lines(normalized_text): normalized_line = normalize_whitespace(raw_line) if not normalized_line: continue upper_line = normalized_line.upper() if upper_line == "BEGIN:VEVENT": if event_properties is None: event_properties = defaultdict(list) continue if upper_line == "END:VEVENT": break parsed = parse_icalendar_property_line(raw_line) if parsed is None: continue target = event_properties if event_properties is not None else calendar_properties target[str(parsed["name"])].append(parsed) if not event_properties: return None def _first_property( properties: dict[str, list[dict[str, object]]], name: str, ) -> dict[str, object] | None: values = properties.get(name) or [] return values[0] if values else None def _first_text( properties: dict[str, list[dict[str, object]]], name: str, ) -> str | None: prop = _first_property(properties, name) if prop is None: return None return normalize_whitespace(str(unescape_icalendar_text(prop.get("value")) or "")) or None summary = normalize_generated_document_title(_first_text(event_properties, "SUMMARY")) description = _first_text(event_properties, "DESCRIPTION") organizer_prop = _first_property(event_properties, "ORGANIZER") organizer = ( format_icalendar_participant(organizer_prop.get("value"), organizer_prop.get("params")) if organizer_prop is not None else None ) attendees: list[str] = [] seen_attendees: set[str] = set() for attendee_prop in event_properties.get("ATTENDEE") or []: formatted = format_icalendar_participant(attendee_prop.get("value"), attendee_prop.get("params")) if not formatted: continue key = formatted.lower() if key in seen_attendees: continue seen_attendees.add(key) attendees.append(formatted) start_prop = _first_property(event_properties, "DTSTART") end_prop = _first_property(event_properties, "DTEND") start_info = ( parse_icalendar_datetime_value( start_prop.get("value"), tzid=(start_prop.get("params") or {}).get("TZID", [None])[0], value_type=(start_prop.get("params") or {}).get("VALUE", [None])[0], ) if start_prop is not None else {} ) end_info = ( parse_icalendar_datetime_value( end_prop.get("value"), tzid=(end_prop.get("params") or {}).get("TZID", [None])[0], value_type=(end_prop.get("params") or {}).get("VALUE", [None])[0], ) if end_prop is not None else {} ) conference_url = ( _first_text(event_properties, "X-GOOGLE-CONFERENCE") or _first_text(event_properties, "URL") or first_icalendar_url(description) ) return { "summary": summary, "description": description, "organizer": organizer, "attendees": attendees, "attendees_display": ", ".join(attendees) or None, "location": _first_text(event_properties, "LOCATION"), "conference_url": conference_url, "start": start_info, "end": end_info, "start_iso": start_info.get("iso"), "end_iso": end_info.get("iso"), "when": format_icalendar_event_range(start_info, end_info), "method": _first_text(calendar_properties, "METHOD"), "status": _first_text(event_properties, "STATUS"), "uid": _first_text(event_properties, "UID"), "sequence": _first_text(event_properties, "SEQUENCE"), } def build_calendar_invite_summary( metadata: dict[str, object] | None, *, file_name: object = None, href: object = None, detail: object = None, ) -> dict[str, str] | None: if not isinstance(metadata, dict): return None title = ( normalize_generated_document_title(metadata.get("summary")) or normalize_generated_document_title(file_name) or "Calendar invite" ) summary = { "kind": "calendar_invite", "label": normalize_whitespace(str(file_name or title or "Calendar invite")) or "Calendar invite", "title": title, "when": normalize_whitespace(str(metadata.get("when") or "")) or "", "organizer": normalize_whitespace(str(metadata.get("organizer") or "")) or "", "attendees": normalize_whitespace(str(metadata.get("attendees_display") or "")) or "", "location": normalize_whitespace(str(metadata.get("location") or "")) or "", "join_href": normalize_whitespace(str(metadata.get("conference_url") or "")) or "", "status": normalize_whitespace(str(summarize_icalendar_invite_status(metadata) or "")) or "", "uid": normalize_whitespace(str(metadata.get("uid") or "")) or "", "sequence": normalize_whitespace(str(metadata.get("sequence") or "")) or "", "href": normalize_whitespace(str(href or "")) or "", "detail": normalize_whitespace(str(detail or "")) or "", "file_name": normalize_whitespace(str(file_name or "")) or "", } if not any(summary.get(key) for key in ("title", "when", "organizer", "attendees", "location", "join_href", "status")): return None return summary def extract_calendar_invite_summary_from_attachment(attachment: dict[str, object]) -> dict[str, str] | None: payload = attachment.get("payload") if not isinstance(payload, (bytes, bytearray)): return None file_name = normalize_whitespace(str(attachment.get("file_name") or "")) or None content_type = normalize_mime_type(attachment.get("content_type")) file_type = infer_attachment_file_type( file_name=file_name, payload=bytes(payload), content_type=content_type, ) if content_type != "text/calendar" and file_type not in ICALENDAR_FILE_TYPES: return None decoded, _, _ = decode_bytes(bytes(payload)) return build_calendar_invite_summary( parse_icalendar_event_metadata(decoded), file_name=file_name, ) def partition_calendar_invite_attachments( attachments: list[dict[str, object]] | None, ) -> tuple[list[dict[str, str]], list[dict[str, object]]]: invite_summaries: list[dict[str, str]] = [] retained_attachments: list[dict[str, object]] = [] for attachment in list(attachments or []): invite_summary = extract_calendar_invite_summary_from_attachment(attachment) if invite_summary is None: retained_attachments.append(attachment) continue invite_summaries.append(invite_summary) return invite_summaries, retained_attachments def build_calendar_invite_search_text(invites: list[dict[str, str]] | None) -> str: if not invites: return "" blocks: list[str] = [] for invite in invites: lines = [CALENDAR_INVITE_TEXT_BLOCK_START] for label, key in ( ("Title", "title"), ("When", "when"), ("Organizer", "organizer"), ("Attendees", "attendees"), ("Location", "location"), ("Join", "join_href"), ("Status", "status"), ("UID", "uid"), ("Sequence", "sequence"), ("Attachment", "file_name"), ): value = normalize_whitespace(str(invite.get(key) or "")) or None if value: lines.append(f"{label}: {value}") lines.append(CALENDAR_INVITE_TEXT_BLOCK_END) blocks.append("\n".join(lines)) return "\n\n".join(blocks) def extract_calendar_invites_from_text_content(text: str) -> tuple[str, list[dict[str, str]]]: invites: list[dict[str, str]] = [] normalized_text = str(text or "").replace("\r\n", "\n").replace("\r", "\n") label_map = { "title": "title", "when": "when", "organizer": "organizer", "attendees": "attendees", "location": "location", "join": "join_href", "status": "status", "uid": "uid", "sequence": "sequence", "attachment": "file_name", } def _replace(match: re.Match[str]) -> str: invite: dict[str, str] = {"kind": "calendar_invite"} for raw_line in match.group(1).splitlines(): if ":" not in raw_line: continue raw_label, raw_value = raw_line.split(":", 1) key = label_map.get(normalize_whitespace(raw_label).lower()) value = normalize_whitespace(raw_value) if key and value: invite[key] = value if invite.get("title") or invite.get("when") or invite.get("join_href"): invite.setdefault("label", invite.get("file_name") or invite.get("title") or "Calendar invite") invites.append(invite) return "" cleaned = CALENDAR_INVITE_TEXT_BLOCK_PATTERN.sub(_replace, normalized_text) cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip("\n") return cleaned, invites CHAT_SPEAKER_BLOCKLIST = { "agenda", "answer", "bcc", "cc", "date", "description", "from", "message", "note", "notes", "owner", "priority", "question", "sent", "status", "subject", "summary", "task", "thread", "title", "to", "topic", } CHAT_TIMESTAMP_HINT_PATTERN = re.compile( r"\d{4}-\d{2}-\d{2}|\d{1,2}/\d{1,2}/\d{2,4}|\b(?:jan|feb|mar|apr|may|jun|jul|aug|sep|sept|oct|nov|dec)[a-z]*\b", re.IGNORECASE, ) CHAT_ISO_DATETIME_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$") CHAT_JSON_FIELD_PATTERN = re.compile(r'^"[^"\n]{1,80}"\s*:\s*') CHAT_LINE_PATTERNS = ( r"^\[(?P<timestamp>[^\]]{4,80})\]\s*(?P<speaker>[^:\n]{2,80}?):\s+(?P<body>\S.*)$", r"^(?P<timestamp>\d{4}-\d{2}-\d{2}[ T]\d{1,2}:\d{2}(?::\d{2})?(?:\s*(?:AM|PM))?(?:\s*(?:Z|UTC|[+\-]\d{2}:?\d{2}))?)\s*[-,]?\s*(?P<speaker>[^:\n]{2,80}?):\s+(?P<body>\S.*)$", r"^(?P<timestamp>\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2}(?::\d{2})?\s*(?:AM|PM)?)\s*[-,]?\s*(?P<speaker>[^:\n]{2,80}?):\s+(?P<body>\S.*)$", r"^(?P<speaker>[^:\n]{2,80}?):\s+(?P<body>\S.*)$", ) def normalize_chat_speaker(value: str | None) -> str | None: candidate = normalize_participant_token(value) if not candidate: return None lowered = candidate.lower().strip("[]()") if lowered in CHAT_SPEAKER_BLOCKLIST or len(candidate.split()) > 8: return None return candidate def parse_chat_timestamp(value: str | None) -> str | None: raw = normalize_whitespace(str(value or "")).strip("[]()") if not raw or not CHAT_TIMESTAMP_HINT_PATTERN.search(raw): return None normalized = normalize_datetime(raw) if normalized and CHAT_ISO_DATETIME_PATTERN.fullmatch(normalized): return normalized return None def format_chat_preview_timestamp(value: object) -> str | None: raw = normalize_whitespace(str(value or "")).strip("[]()") if not raw: return None parsed = parse_utc_timestamp(raw) if parsed is None: normalized = parse_chat_timestamp(raw) parsed = parse_utc_timestamp(normalized) if normalized else None if parsed is None: return raw return parsed.strftime("%b %d, %Y %I:%M %p UTC").replace(" 0", " ") def chat_avatar_initials(value: str) -> str: letters = [part[0].upper() for part in re.split(r"\s+", value.strip()) if part and part[0].isalnum()] if not letters: return "?" if len(letters) == 1: return letters[0] return f"{letters[0]}{letters[-1]}" CHAT_AVATAR_PALETTE = ( ("#dbeafe", "#1d4ed8"), ("#dcfce7", "#166534"), ("#fef3c7", "#92400e"), ("#fce7f3", "#9d174d"), ("#ede9fe", "#6d28d9"), ("#cffafe", "#0f766e"), ("#fee2e2", "#b91c1c"), ("#e0e7ff", "#4338ca"), ) def normalize_chat_avatar_color(value: object) -> str | None: candidate = normalize_whitespace(str(value or "")).lstrip("#") if re.fullmatch(r"[0-9a-fA-F]{6}", candidate or ""): return f"#{candidate.lower()}" return None def chat_avatar_colors(seed: str, preferred_background: object = None) -> tuple[str, str]: background = normalize_chat_avatar_color(preferred_background) if background: red = int(background[1:3], 16) green = int(background[3:5], 16) blue = int(background[5:7], 16) luminance = (0.2126 * red) + (0.7152 * green) + (0.0722 * blue) return background, ("#ffffff" if luminance < 140 else "#111827") palette_index = int(hashlib.sha1(seed.encode("utf-8")).hexdigest(), 16) % len(CHAT_AVATAR_PALETTE) return CHAT_AVATAR_PALETTE[palette_index] def build_chat_avatar_svg(label: str, background: str, foreground: str, alt_text: str) -> str: return ( '<svg class="chat-avatar-svg" xmlns="http://www.w3.org/2000/svg" width="96" height="96" ' 'viewBox="0 0 96 96" role="img" ' f'aria-label="{html.escape(alt_text, quote=True)}">' f'<circle cx="48" cy="48" r="48" fill="{html.escape(background, quote=True)}"/>' f'<text x="50%" y="55%" text-anchor="middle" dominant-baseline="middle" ' f'font-family="Inter, -apple-system, BlinkMacSystemFont, Segoe UI, sans-serif" ' f'font-size="30" font-weight="700" fill="{html.escape(foreground, quote=True)}">{html.escape(label)}</text>' "</svg>" ) def iter_chat_transcript_entries(text: str, max_lines: int = 800) -> list[dict[str, str | None]]: if not text: return [] entries: list[dict[str, str | None]] = [] for raw_line in text.splitlines()[:max_lines]: stripped = raw_line.strip() if not stripped or len(stripped) > 240: continue if CHAT_JSON_FIELD_PATTERN.match(stripped): continue for pattern in CHAT_LINE_PATTERNS: match = re.match(pattern, stripped, flags=re.IGNORECASE) if not match: continue speaker = normalize_chat_speaker(match.groupdict().get("speaker")) body = normalize_whitespace(match.groupdict().get("body") or "") if not speaker or not body: continue entries.append( { "speaker": speaker, "body": body, "timestamp": parse_chat_timestamp(match.groupdict().get("timestamp")), } ) break return entries def extract_chat_participants(text: str) -> str | None: participants: list[str] = [] seen: set[str] = set() speaker_counts: dict[str, int] = {} timestamped_matches = 0 for entry in iter_chat_transcript_entries(text): candidate = str(entry["speaker"]) key = candidate.lower().strip("[]()") speaker_counts[key] = speaker_counts.get(key, 0) + 1 if isinstance(entry.get("timestamp"), str): timestamped_matches += 1 if key not in seen: seen.add(key) participants.append(candidate) total_matches = sum(speaker_counts.values()) if total_matches < 2: return None if timestamped_matches < 2: repeated_speaker = any(count >= 2 for count in speaker_counts.values()) if total_matches < 3 or not repeated_speaker: return None elif len(participants) < 2 and total_matches < 3: return None return ", ".join(participants) def extract_chat_transcript_metadata(text: str) -> dict[str, object] | None: entries = iter_chat_transcript_entries(text, max_lines=1200) if not entries: return None participants: list[str] = [] seen: set[str] = set() speaker_counts: dict[str, int] = {} first_speaker: str | None = None first_body: str | None = None first_timestamp: str | None = None last_timestamp: str | None = None timestamped_matches = 0 for entry in entries: speaker = str(entry["speaker"]) key = speaker.lower().strip("[]()") speaker_counts[key] = speaker_counts.get(key, 0) + 1 if key not in seen: seen.add(key) participants.append(speaker) if first_speaker is None: first_speaker = speaker if first_body is None: first_body = str(entry["body"]) timestamp = entry.get("timestamp") if isinstance(timestamp, str): timestamped_matches += 1 if first_timestamp is None: first_timestamp = timestamp last_timestamp = timestamp total_matches = sum(speaker_counts.values()) repeated_speaker = any(count >= 2 for count in speaker_counts.values()) if total_matches < 2: return None if timestamped_matches < 2: if len(participants) < 2 or total_matches < 3 or not repeated_speaker: return None elif len(participants) < 2 and total_matches < 3: return None return { "author": first_speaker, "participants": ", ".join(participants) or None, "date_created": first_timestamp, "date_modified": last_timestamp if last_timestamp and last_timestamp != first_timestamp else None, "title": (first_body[:200] if first_body else None), "message_count": total_matches, "timestamped_message_count": timestamped_matches, } def infer_content_type_from_content( file_type: str, text_content: str, email_headers: dict[str, str | None] | None = None, chat_metadata: dict[str, object] | None = None, ) -> str | None: if email_headers: return "Email" if chat_metadata: return "Chat" if not text_content: return None leading_text = text_content[:4000].upper() if "BEGIN:VCALENDAR" in leading_text or "BEGIN:VEVENT" in leading_text: return "Calendar" if file_type in {"xml"} and "<VCALENDAR" in leading_text: return "Calendar" return None def determine_content_type( path: Path, text_content: str, email_headers: dict[str, str | None] | None = None, chat_metadata: dict[str, object] | None = None, explicit_content_type: str | None = None, ) -> str | None: file_type = normalize_extension(path) return ( infer_content_type_from_content(file_type, text_content, email_headers, chat_metadata) or explicit_content_type or infer_content_type_from_extension(file_type) ) LAZY_DEPENDENCY_IMPORT_TARGETS = { "charset_normalizer": ("charset_normalizer", None), "extract_msg": ("extract_msg", None), "openpyxl": ("openpyxl", None), "xlrd": ("xlrd", None), "pdfplumber": ("pdfplumber", None), "DocxDocument": ("docx", "Document"), "rtf_to_text": ("striprtf.striprtf", "rtf_to_text"), "PilImage": ("PIL.Image", None), "pypff": ("pypff", None), } def import_dependency_target(module_name: str, attribute_name: str | None) -> object: imported = importlib.import_module(module_name) return imported if attribute_name is None else getattr(imported, attribute_name) def load_dependency(dependency_name: str, *, allow_auto_install: bool = True) -> object | None: current = globals().get(dependency_name, _UNLOADED_DEPENDENCY) if current is not _UNLOADED_DEPENDENCY and current is not None: return current import_target = LAZY_DEPENDENCY_IMPORT_TARGETS.get(dependency_name) if import_target is None: raise RetrieverError(f"Unknown dependency loader: {dependency_name}") module_name, attribute_name = import_target try: value = import_dependency_target(module_name, attribute_name) except Exception: value = None runtime_paths = plugin_runtime_paths(root=ACTIVE_WORKSPACE_ROOT) if value is None and runtime_paths is not None: try: if activate_plugin_site_packages(runtime_paths): value = import_dependency_target(module_name, attribute_name) except Exception: value = None if value is None and allow_auto_install and runtime_paths is not None: try: ensure_plugin_runtime( runtime_paths, install_requirements=True, force_requirements_install=True, reason=f"dependency:{dependency_name}", ) activate_plugin_site_packages(runtime_paths) value = import_dependency_target(module_name, attribute_name) except Exception: value = None globals()[dependency_name] = value return value def dependency_status( dependency_name: str, *, package_name: str, import_name: str | None = None, detail_label: str | None = None, probe_if_unloaded: bool = False, allow_auto_install: bool = False, ) -> dict[str, str]: current = globals().get(dependency_name, _UNLOADED_DEPENDENCY) if current is _UNLOADED_DEPENDENCY and probe_if_unloaded: current = load_dependency(dependency_name, allow_auto_install=allow_auto_install) detail_name = detail_label or import_name or dependency_name import_label = import_name or dependency_name if current is _UNLOADED_DEPENDENCY: return { "status": "deferred", "detail": f"{detail_name} will load on demand when a matching file type is used.", } if current is None: return { "status": "fail", "detail": f"Missing optional dependency import '{import_label}'. Install {package_name} before using the matching file type.", } return {"status": "pass", "detail": f"{detail_name} import succeeded"} def dependency_guard(module: object | str | None, package_name: str, file_type: str) -> object: if isinstance(module, str): module = load_dependency(module, allow_auto_install=True) if module is None: raise RetrieverError( f"Missing dependency for .{file_type} parsing: install {package_name} before ingesting this file type." ) if module is _UNLOADED_DEPENDENCY: raise RetrieverError( f"Dependency for .{file_type} parsing was not initialized correctly: {package_name}." ) return module CID_REFERENCE_PATTERN = re.compile( r"""(?i)(\b(?:src|background)\s*=\s*)(["'])cid:([^"']+)\2""" ) def sniff_image_mime_type(payload: bytes) -> str | None: if not isinstance(payload, (bytes, bytearray)) or len(payload) < 4: return None data = bytes(payload) if data.startswith(b"\x89PNG\r\n\x1a\n"): return "image/png" if data.startswith(b"\xff\xd8\xff"): return "image/jpeg" if data.startswith(b"GIF87a") or data.startswith(b"GIF89a"): return "image/gif" if data.startswith(b"BM"): return "image/bmp" if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP": return "image/webp" if data.startswith(b"II*\x00") or data.startswith(b"MM\x00*"): return "image/tiff" return None def decode_attachment_text_sample(payload: bytes, *, max_bytes: int = 65536) -> str | None: if not isinstance(payload, (bytes, bytearray)): return None sample = bytes(payload[:max_bytes]) if not sample or b"\x00" in sample: return None decoded, _, _ = decode_bytes(sample) if not decoded: return None replacement_count = decoded.count("\ufffd") if replacement_count > max(4, len(decoded) // 50): return None control_count = sum( 1 for character in decoded if ord(character) < 32 and character not in "\r\n\t\f\b" ) if control_count > max(4, len(decoded) // 50): return None return decoded def sniff_attachment_file_type(payload: bytes) -> str | None: if not isinstance(payload, (bytes, bytearray)) or not payload: return None data = bytes(payload) sample = data[:65536] trimmed_sample = sample.lstrip(b"\xef\xbb\xbf\r\n\t ") if trimmed_sample.startswith(b"%PDF-"): return "pdf" image_mime_type = sniff_image_mime_type(data) if image_mime_type: return attachment_file_type_from_mime_type(image_mime_type) if trimmed_sample.startswith(b"{\\rtf"): return "rtf" if zipfile.is_zipfile(io.BytesIO(data)): try: with zipfile.ZipFile(io.BytesIO(data)) as archive: member_names = set(archive.namelist()) except Exception: member_names = set() if "[Content_Types].xml" in member_names: if any(name.startswith("word/") for name in member_names): return "docx" if any(name.startswith("xl/") for name in member_names): return "xlsx" if any(name.startswith("ppt/") for name in member_names): return "pptx" return "zip" if sample.startswith(OLE_COMPOUND_FILE_MAGIC): if b"Workbook" in sample or b"Book" in sample: return "xls" if b"WordDocument" in sample: return "doc" if b"PowerPoint Document" in sample: return "ppt" return "ole" decoded_text = decode_attachment_text_sample(data) if not decoded_text: return None stripped_text = decoded_text.lstrip("\ufeff\r\n\t ") if not stripped_text: return None preview = stripped_text[:1024] preview_lower = preview.lower() if stripped_text.upper().startswith("BEGIN:VCALENDAR"): return "ics" if ( preview_lower.startswith("<!doctype html") or preview_lower.startswith("<html") or preview_lower.startswith("<body") or "<html" in preview_lower ): return "html" if stripped_text.startswith("{") or stripped_text.startswith("["): try: parsed = json.loads(stripped_text) except Exception: parsed = None if isinstance(parsed, (dict, list)): return "json" if stripped_text.startswith("<?xml") or stripped_text.startswith("<"): try: ET.fromstring(stripped_text) except Exception: pass else: return "xml" return "txt" def infer_attachment_file_type( *, file_name: str | None = None, payload: bytes | None = None, content_type: object = None, preferred_extension: object = None, ) -> str | None: normalized_extension = normalize_file_type_name(preferred_extension) if normalized_extension: return normalized_extension sniffed = sniff_attachment_file_type(payload) if isinstance(payload, (bytes, bytearray)) else None if sniffed: return sniffed declared = attachment_file_type_from_mime_type(content_type) if declared: return declared if file_name: return normalize_file_type_name(Path(file_name).suffix.lower().lstrip(".")) return None def normalize_content_id(raw: object) -> str | None: if raw is None: return None if isinstance(raw, (bytes, bytearray)): try: raw = bytes(raw).decode("utf-8") except Exception: raw = bytes(raw).decode("utf-8", errors="replace") value = str(raw).strip() if not value: return None value = value.strip("<>").strip() return value or None def attachment_image_mime_type(attachment: object) -> str | None: if not isinstance(attachment, dict): return None payload = attachment.get("payload") payload_bytes = bytes(payload) if isinstance(payload, (bytes, bytearray)) else None file_name = str(attachment.get("file_name") or "") mime_type = normalize_mime_type(attachment.get("content_type")) if mime_type is not None and mime_type.startswith("image/"): return mime_type ooxml_mime_type = ooxml_image_mime_type(file_name) if ooxml_mime_type: return ooxml_mime_type if payload_bytes is not None: sniffed = sniff_image_mime_type(payload_bytes) if sniffed: return sniffed guessed, _ = mimetypes.guess_type(file_name) if guessed and guessed.startswith("image/"): return guessed return None def build_cid_data_uri_map(attachments: list[dict[str, object]] | None) -> dict[str, str]: if not attachments: return {} mapping: dict[str, str] = {} for attachment in attachments: if not isinstance(attachment, dict): continue content_id = normalize_content_id(attachment.get("content_id")) if not content_id: continue payload = attachment.get("payload") if not isinstance(payload, (bytes, bytearray)): continue payload_bytes = bytes(payload) mime_type = attachment_image_mime_type(attachment) if not mime_type: mime_type = "application/octet-stream" encoded = base64.b64encode(payload_bytes).decode("ascii") mapping[content_id.lower()] = f"data:{mime_type};base64,{encoded}" return mapping def inline_cid_references_in_html( html_body: str | None, attachments: list[dict[str, object]] | None, ) -> str | None: if not html_body: return html_body cid_map = build_cid_data_uri_map(attachments) if not cid_map: return html_body def _replace(match: re.Match[str]) -> str: prefix = match.group(1) quote = match.group(2) cid = normalize_content_id(match.group(3)) if not cid: return match.group(0) replacement = cid_map.get(cid.lower()) if not replacement: return match.group(0) return f"{prefix}{quote}{replacement}{quote}" return CID_REFERENCE_PATTERN.sub(_replace, html_body) def referenced_cids_in_html(html_body: str | None) -> set[str]: if not html_body: return set() referenced: set[str] = set() for match in CID_REFERENCE_PATTERN.finditer(html_body): cid = normalize_content_id(match.group(3)) if cid: referenced.add(cid.lower()) return referenced def filter_html_preview_embedded_image_attachments( html_body: str | None, attachments: list[dict[str, object]] | None, ) -> list[dict[str, object]]: if not attachments: return [] referenced_cids = referenced_cids_in_html(html_body) if not referenced_cids: return list(attachments) filtered: list[dict[str, object]] = [] for attachment in attachments: content_id = normalize_content_id(attachment.get("content_id")) if content_id and content_id.lower() in referenced_cids and attachment_image_mime_type(attachment): continue filtered.append(attachment) return filtered def render_html_preview_calendar_invite_cards(links: list[dict[str, str]]) -> str: if not links: return "" cards: list[str] = [] for link in links: title = normalize_whitespace(str(link.get("title") or link.get("label") or "Calendar invite")) or "Calendar invite" href = normalize_whitespace(str(link.get("href") or "")) or None title_html = ( f'<a href="{html.escape(href)}">{html.escape(title)}</a>' if href else html.escape(title) ) metadata_items: list[str] = [] for label, key in ( ("When", "when"), ("Organizer", "organizer"), ("Attendees", "attendees"), ("Location", "location"), ("Status", "status"), ): value = normalize_whitespace(str(link.get(key) or "")) if not value: continue metadata_items.append( f"<div><dt>{html.escape(label)}</dt><dd>{html.escape(value)}</dd></div>" ) join_href = normalize_whitespace(str(link.get("join_href") or "")) or None if join_href: metadata_items.append( "<div><dt>Join</dt>" f'<dd><a href="{html.escape(join_href)}">{html.escape(join_href)}</a></dd></div>' ) detail = normalize_whitespace(str(link.get("detail") or "")) detail_html = ( f'<p class="retriever-calendar-invite-detail">{html.escape(detail)}</p>' if detail else "" ) cards.append( '<article class="retriever-calendar-invite">' '<div class="retriever-calendar-invite-header">' "<div>" '<p class="retriever-calendar-invite-kicker">Calendar invite</p>' f'<h3 class="retriever-calendar-invite-title">{title_html}</h3>' "</div>" f"{detail_html}" "</div>" + ( f'<dl class="retriever-calendar-invite-meta">{"".join(metadata_items)}</dl>' if metadata_items else "" ) + "</article>" ) return ( "<!-- RETRIEVER_CALENDAR_INVITES_START -->" + '<section class="retriever-calendar-invites">' + "".join(cards) + "</section>" + "<!-- RETRIEVER_CALENDAR_INVITES_END -->" ) def render_html_preview_attachment_links(links: list[dict[str, str]]) -> str: if not links: return "" calendar_links = [ link for link in links if normalize_whitespace(str(link.get("kind") or "")).lower() == "calendar_invite" ] file_links = [ link for link in links if normalize_whitespace(str(link.get("kind") or "")).lower() != "calendar_invite" ] sections: list[str] = [] calendar_section = render_html_preview_calendar_invite_cards(calendar_links) if calendar_section: sections.append(calendar_section) items: list[str] = [] for link in file_links: href = html.escape(str(link.get("href") or "")) label = html.escape(str(link.get("label") or "Attachment")) detail = normalize_whitespace(str(link.get("detail") or "")) detail_html = f' <span class="retriever-attachment-meta">({html.escape(detail)})</span>' if detail else "" items.append(f'<li><a href="{href}">{label}</a>{detail_html}</li>') if items: sections.append( '<section class="retriever-attachments"><h2>Attachments</h2><ul>' + "".join(items) + "</ul></section>" ) if not sections: return "" return ( "<!-- RETRIEVER_ATTACHMENT_LINKS_START -->" + "".join(sections) + "<!-- RETRIEVER_ATTACHMENT_LINKS_END -->" ) def inject_html_preview_attachment_links(html_text: str, links: list[dict[str, str]]) -> str: cleaned = HTML_PREVIEW_ATTACHMENT_LINKS_PATTERN.sub("", html_text) section = render_html_preview_attachment_links(links) if not section: return cleaned if "</h1>" in cleaned: return cleaned.replace("</h1>", f"</h1>{section}", 1) if "<body>" in cleaned: return cleaned.replace("<body>", f"<body>{section}", 1) return cleaned + section def build_html_preview( headers: dict[str, str], body_html: str | None = None, body_text: str | None = None, *, document_title: str, head_html: str | None = None, heading: str | None = None, ) -> str: header_html = "".join( f"<tr><th>{html.escape(key)}</th><td>{html.escape(value)}</td></tr>" for key, value in headers.items() if value ) resolved_heading = document_title if heading is None else heading heading_html = f"<h1>{html.escape(resolved_heading)}</h1>" if resolved_heading else "" header_section = f"<table>{header_html}</table><hr/>" if header_html else "" if body_html: body_section = body_html else: body_section = f"<pre>{html.escape(body_text or '')}</pre>" return ( "<!DOCTYPE html>" "<html><head>" '<meta charset="utf-8"/>' '<meta name="viewport" content="width=device-width, initial-scale=1"/>' f"<title>{html.escape(document_title)}" f"{head_html or ''}" "" f"{heading_html}" f"{header_section}" f"{body_section}" "" ) def build_chat_preview_html( headers: dict[str, str], body_text: str, *, document_title: str, entries: list[dict[str, object]] | None = None, ) -> str: chat_entries = entries if entries is not None else iter_chat_transcript_entries(body_text, max_lines=4000) head_html = ( "" ) if chat_entries: rendered_entries: list[str] = [] for entry in chat_entries: speaker = normalize_whitespace(str(entry.get("speaker") or "")) or "Unknown" body = str(entry.get("body") or "").strip() if not body: continue timestamp_label = ( normalize_whitespace(str(entry.get("timestamp_label") or "")) or format_chat_preview_timestamp(entry.get("timestamp")) or "" ) timestamp_html = f'[{html.escape(timestamp_label)}]' if timestamp_label else "" avatar_label = normalize_whitespace(str(entry.get("avatar_label") or "")) or chat_avatar_initials(speaker) avatar_background, avatar_foreground = chat_avatar_colors( speaker, entry.get("avatar_color"), ) avatar_html = build_chat_avatar_svg(avatar_label, avatar_background, avatar_foreground, speaker) rendered_entries.append( "
" f"{avatar_html}" "
" "
" f"{html.escape(speaker)}" f"{timestamp_html}" "
" f"
{html.escape(body)}
" "
" "
" ) if rendered_entries: body_section = ( "
" f"{''.join(rendered_entries)}" "
" "
" "Full transcript" f"
{html.escape(body_text or '')}
" "
" ) else: body_section = f"
{html.escape(body_text or '')}
" else: body_section = f"
{html.escape(body_text or '')}
" return build_html_preview( headers, body_html=body_section, document_title=document_title, head_html=head_html, ) def conversation_preview_anchor(document_id: int) -> str: return f"doc-{int(document_id)}" def conversation_preview_base_path(conversation_id: int) -> Path: return Path("previews") / "conversations" / f"conversation-{int(conversation_id):08d}" def conversation_preview_full_rel_path(conversation_id: int) -> str: return (conversation_preview_base_path(conversation_id) / "conversation.html").as_posix() def conversation_preview_toc_rel_path(conversation_id: int) -> str: return (conversation_preview_base_path(conversation_id) / "index.html").as_posix() def conversation_preview_segment_rel_path(conversation_id: int, segment_token: str) -> str: normalized_token = re.sub(r"[^A-Za-z0-9._-]+", "-", normalize_whitespace(segment_token) or "segment").strip("-") return (conversation_preview_base_path(conversation_id) / f"segment-{normalized_token or 'segment'}.html").as_posix() def conversation_preview_entry_rel_path(conversation_id: int, document_id: int) -> str: return (conversation_preview_base_path(conversation_id) / f"{conversation_preview_anchor(document_id)}.html").as_posix() def is_conversation_preview_rel_path(rel_preview_path: object) -> bool: normalized = normalize_internal_rel_path(Path(str(rel_preview_path or ""))) return normalized.startswith("previews/conversations/") def append_preview_fragment(path: str, target_fragment: object) -> str: fragment = normalize_whitespace(str(target_fragment or "")) if not fragment: return path if "#" in path: return path return f"{path}#{fragment}" def parse_xml_document(data: bytes) -> ET.Element: return ET.fromstring(data) def ooxml_relationship_part_name(part_name: str) -> str: directory, file_name = posixpath.split(part_name) return posixpath.join(directory, "_rels", f"{file_name}.rels") def normalize_ooxml_target(base_part: str, target: str) -> str: return posixpath.normpath(posixpath.join(posixpath.dirname(base_part), target)) def read_ooxml_relationships(archive: zipfile.ZipFile, part_name: str) -> dict[str, dict[str, str]]: rels_part = ooxml_relationship_part_name(part_name) try: root = parse_xml_document(archive.read(rels_part)) except KeyError: return {} relationships: dict[str, dict[str, str]] = {} for relationship in root.findall("rels:Relationship", OOXML_RELATIONSHIP_NS): rel_id = relationship.attrib.get("Id") target = relationship.attrib.get("Target") rel_type = relationship.attrib.get("Type") if rel_id and target and rel_type: relationships[rel_id] = { "target": normalize_ooxml_target(part_name, target), "type": rel_type, } return relationships def xml_local_name(tag: str) -> str: return tag.split("}", 1)[1] if "}" in tag else tag def pptx_shape_position(element: ET.Element) -> tuple[int, int]: for query in ("./p:spPr/a:xfrm/a:off", "./p:xfrm/a:off", "./p:grpSpPr/a:xfrm/a:off"): offset = element.find(query, PPTX_NAMESPACES) if offset is not None: x = int(offset.attrib.get("x", "0") or "0") y = int(offset.attrib.get("y", "0") or "0") return x, y return 0, 0 def pptx_shape_size(element: ET.Element) -> tuple[int, int]: for query in ("./p:spPr/a:xfrm/a:ext", "./p:xfrm/a:ext", "./p:grpSpPr/a:xfrm/a:ext"): extent = element.find(query, PPTX_NAMESPACES) if extent is not None: cx = int(extent.attrib.get("cx", "0") or "0") cy = int(extent.attrib.get("cy", "0") or "0") return cx, cy return 0, 0 def pptx_shape_placeholder_type(element: ET.Element) -> str | None: for query in ( "./p:nvSpPr/p:nvPr/p:ph", "./p:nvGraphicFramePr/p:nvPr/p:ph", "./p:nvGrpSpPr/p:nvPr/p:ph", ): placeholder = element.find(query, PPTX_NAMESPACES) if placeholder is not None: return placeholder.attrib.get("type") or "body" return None def pptx_paragraph_text(paragraph: ET.Element) -> str: parts = [text_node.text or "" for text_node in paragraph.findall(".//a:t", PPTX_NAMESPACES)] return normalize_whitespace("".join(parts)) def ooxml_image_mime_type(part_name: str) -> str | None: suffix = Path(part_name).suffix.lower() if suffix == ".png": return "image/png" if suffix in {".jpg", ".jpeg"}: return "image/jpeg" if suffix == ".gif": return "image/gif" if suffix == ".bmp": return "image/bmp" if suffix == ".webp": return "image/webp" if suffix in {".tif", ".tiff"}: return "image/tiff" return None _PREVIEW_IMAGE_OUTPUT_PROFILE: tuple[str, str, str] | None = None def preview_image_output_profile() -> tuple[str, str, str]: global _PREVIEW_IMAGE_OUTPUT_PROFILE if _PREVIEW_IMAGE_OUTPUT_PROFILE is not None: return _PREVIEW_IMAGE_OUTPUT_PROFILE pil_image_module = load_dependency("PilImage") if pil_image_module is None: _PREVIEW_IMAGE_OUTPUT_PROFILE = ("PNG", "image/png", ".png") return _PREVIEW_IMAGE_OUTPUT_PROFILE try: probe_buffer = io.BytesIO() pil_image_module.new("L", (1, 1), 255).save(probe_buffer, format="WEBP", lossless=True) _PREVIEW_IMAGE_OUTPUT_PROFILE = ("WEBP", "image/webp", ".webp") except Exception: _PREVIEW_IMAGE_OUTPUT_PROFILE = ("PNG", "image/png", ".png") return _PREVIEW_IMAGE_OUTPUT_PROFILE def preview_image_output_mime_type() -> str: return preview_image_output_profile()[1] def preview_image_output_suffix() -> str: return preview_image_output_profile()[2] def image_path_preview_raster(path: Path, *, max_dimension: int | None = None) -> tuple[bytes, str, str] | None: resized_dimension = max(0, int(max_dimension or 0)) pil_image_module = load_dependency("PilImage") if pil_image_module is None: return None output_format, mime_type, output_suffix = preview_image_output_profile() with pil_image_module.open(path) as image: restore_bilevel = False if image.mode == "1": if output_format == "WEBP": # Lossless WebP compresses ordinary bilevel production scans # better at source resolution than after an antialiased resize. image = image.convert("L") resized_dimension = 0 elif resized_dimension: image = image.convert("L") restore_bilevel = True elif image.mode == "P": if output_format == "WEBP" or resized_dimension: image = image.convert("RGBA" if image.info.get("transparency") is not None else "RGB") elif output_format == "WEBP" and image.mode not in {"L", "RGB", "RGBA"}: image = image.convert("RGBA" if "A" in image.getbands() else "RGB") if resized_dimension: # Normalize low-bit-depth images before resizing so Pillow can use # antialiased resampling instead of mode-1 nearest-neighbor. if image.mode == "1": image = image.convert("L") restore_bilevel = True resampling = getattr(pil_image_module, "Resampling", pil_image_module) lanczos = getattr(resampling, "LANCZOS", None) if lanczos is None: image.thumbnail((resized_dimension, resized_dimension)) else: try: image.thumbnail((resized_dimension, resized_dimension), resample=lanczos) except ValueError: image.thumbnail((resized_dimension, resized_dimension)) if restore_bilevel: # Threshold back to bilevel after the antialiased resize so large # embedded production previews stay compact. dither = getattr(getattr(pil_image_module, "Dither", pil_image_module), "NONE", None) image = image.convert("1") if dither is None else image.convert("1", dither=dither) buffer = io.BytesIO() if output_format == "WEBP": try: image.save(buffer, format="WEBP", lossless=True) return buffer.getvalue(), mime_type, output_suffix except (OSError, ValueError): output_format = "PNG" mime_type = "image/png" output_suffix = ".png" buffer = io.BytesIO() try: image.save(buffer, format="PNG", optimize=True) except (OSError, ValueError): image.convert("RGBA" if "A" in image.getbands() else "RGB").save(buffer, format="PNG", optimize=True) return buffer.getvalue(), mime_type, output_suffix def image_path_data_url(path: Path, *, max_dimension: int | None = None) -> str | None: mime_type, _ = mimetypes.guess_type(path.name) normalized_suffix = path.suffix.lower() resized_dimension = max(0, int(max_dimension or 0)) if normalized_suffix in {".tif", ".tiff"} or resized_dimension: preview_raster = image_path_preview_raster(path, max_dimension=resized_dimension) if preview_raster is None: if normalized_suffix not in {".tif", ".tiff"} and mime_type is not None and mime_type.startswith("image/"): return f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('ascii')}" return None preview_bytes, preview_mime_type, _ = preview_raster return f"data:{preview_mime_type};base64,{base64.b64encode(preview_bytes).decode('ascii')}" if mime_type is None or not mime_type.startswith("image/"): return None return f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('ascii')}" READ_SAFE_VISUAL_SUFFIXES = frozenset( { ".pdf", ".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", } ) def ensure_read_safe_visual_artifact_path(root: Path, artifact_path: Path) -> Path: if artifact_path.suffix.lower() in READ_SAFE_VISUAL_SUFFIXES: return artifact_path pil_image_module = load_dependency("PilImage") if pil_image_module is None: raise RetrieverError( f"Could not load Pillow to convert unsupported visual artifact {artifact_path.name!r} into a Read-safe PNG." ) paths = workspace_paths(root) output_dir = Path(paths["tmp_dir"]) / "read-safe-visual-artifacts" output_dir.mkdir(parents=True, exist_ok=True) stat_result = artifact_path.stat() fingerprint = hashlib.sha256( f"{artifact_path.resolve()}:{stat_result.st_size}:{stat_result.st_mtime_ns}".encode("utf-8") ).hexdigest()[:12] target_path = output_dir / f"{artifact_path.stem}-{fingerprint}.png" if target_path.exists(): return target_path with pil_image_module.open(artifact_path) as source_image: converted = source_image.convert("RGBA" if "A" in source_image.getbands() else "RGB") converted.save(target_path, format="PNG") return target_path def pptx_picture_entry( element: ET.Element, *, archive: zipfile.ZipFile, relationships: dict[str, dict[str, str]], ) -> dict[str, object] | None: blip = element.find(".//a:blip", PPTX_NAMESPACES) relationship_id = blip.attrib.get(f"{{{PPTX_NAMESPACES['r']}}}embed") if blip is not None else None if not relationship_id: return None relationship = relationships.get(relationship_id) if relationship is None: return None target = relationship["target"] mime_type = ooxml_image_mime_type(target) if mime_type is None: return None try: image_bytes = archive.read(target) except KeyError: return None c_nv_pr = element.find("./p:nvPicPr/p:cNvPr", PPTX_NAMESPACES) alt_text = None if c_nv_pr is not None: alt_text = normalize_whitespace(c_nv_pr.attrib.get("descr", "") or c_nv_pr.attrib.get("name", "")) width_emu, height_emu = pptx_shape_size(element) return { "kind": "image", "alt": alt_text or Path(target).name, "src": f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('ascii')}", "width_px": max(1, round(width_emu / EMU_PER_PIXEL)) if width_emu else None, "height_px": max(1, round(height_emu / EMU_PER_PIXEL)) if height_emu else None, } def pptx_shape_text_blocks(element: ET.Element) -> list[str]: local_name = xml_local_name(element.tag) if local_name == "sp": paragraphs = [ text for paragraph in element.findall("./p:txBody/a:p", PPTX_NAMESPACES) if (text := pptx_paragraph_text(paragraph)) ] return ["\n".join(paragraphs)] if paragraphs else [] if local_name == "graphicFrame": table = element.find(".//a:tbl", PPTX_NAMESPACES) if table is None: return [] rows: list[str] = [] for row in table.findall("./a:tr", PPTX_NAMESPACES): cells = [ text for cell in row.findall("./a:tc", PPTX_NAMESPACES) if (text := normalize_whitespace(" ".join(filter(None, [node.text for node in cell.findall('.//a:t', PPTX_NAMESPACES)])))) ] if cells: rows.append(" | ".join(cells)) return rows return [] def collect_pptx_shape_entries( container: ET.Element, *, archive: zipfile.ZipFile | None = None, relationships: dict[str, dict[str, str]] | None = None, group_offset: tuple[int, int] = (0, 0), ) -> list[dict[str, object]]: entries: list[dict[str, object]] = [] sequence = 0 for child in list(container): local_name = xml_local_name(child.tag) if local_name in {"nvGrpSpPr", "grpSpPr"}: continue if local_name == "grpSp": child_x, child_y = pptx_shape_position(child) entries.extend( collect_pptx_shape_entries( child, archive=archive, relationships=relationships, group_offset=(group_offset[0] + child_x, group_offset[1] + child_y), ) ) sequence += 1 continue if local_name == "pic" and archive is not None and relationships is not None: image_entry = pptx_picture_entry(child, archive=archive, relationships=relationships) if image_entry is not None: child_x, child_y = pptx_shape_position(child) entries.append( { **image_entry, "placeholder_type": None, "x": group_offset[0] + child_x, "y": group_offset[1] + child_y, "sequence": sequence, } ) sequence += 1 continue text_blocks = pptx_shape_text_blocks(child) if not text_blocks: sequence += 1 continue child_x, child_y = pptx_shape_position(child) placeholder_type = pptx_shape_placeholder_type(child) entries.append( { "kind": "text", "blocks": text_blocks, "placeholder_type": placeholder_type, "x": group_offset[0] + child_x, "y": group_offset[1] + child_y, "sequence": sequence, } ) sequence += 1 return entries def sorted_pptx_content_entries( container: ET.Element, *, archive: zipfile.ZipFile | None = None, relationships: dict[str, dict[str, str]] | None = None, ) -> list[dict[str, object]]: entries = collect_pptx_shape_entries(container, archive=archive, relationships=relationships) return sorted( entries, key=lambda item: ( 0 if item["placeholder_type"] in {"title", "ctrTitle", "subTitle"} else 1, int(item["y"]), int(item["x"]), int(item["sequence"]), ), ) def sorted_pptx_text_blocks(container: ET.Element) -> list[str]: ordered = sorted_pptx_content_entries(container) blocks: list[str] = [] for entry in ordered: if entry.get("kind") == "text": blocks.extend(str(block) for block in entry["blocks"]) return blocks def render_html_text_blocks(blocks: list[str]) -> str: if not blocks: return "

No extractable text.

" paragraphs = [] for block in blocks: escaped = html.escape(block).replace("\n", "
") paragraphs.append(f"

{escaped}

") return "".join(paragraphs) def render_pptx_content_entries(entries: list[dict[str, object]]) -> str: if not entries: return "

No extractable content.

" rendered: list[str] = [] for entry in entries: if entry.get("kind") == "image": alt_text = str(entry.get("alt") or "Slide image") width_px = entry.get("width_px") height_px = entry.get("height_px") size_attrs = "" if isinstance(width_px, int) and width_px > 0: size_attrs += f' width="{width_px}"' if isinstance(height_px, int) and height_px > 0: size_attrs += f' height="{height_px}"' rendered.append( '
' f'{html.escape(alt_text)}' f"
{html.escape(alt_text)}
" "
" ) continue rendered.append(render_html_text_blocks([str(block) for block in entry.get("blocks", [])])) return "".join(rendered) def extract_pptx_notes_blocks(archive: zipfile.ZipFile, slide_part_name: str) -> list[str]: relationships = read_ooxml_relationships(archive, slide_part_name) notes_part_name = None for relationship in relationships.values(): if relationship["type"] == PPTX_NOTES_RELATIONSHIP_TYPE: notes_part_name = relationship["target"] break if not notes_part_name: return [] try: notes_root = parse_xml_document(archive.read(notes_part_name)) except KeyError: return [] notes_tree = notes_root.find("./p:cSld/p:spTree", PPTX_NAMESPACES) if notes_tree is None: return [] return sorted_pptx_text_blocks(notes_tree) def build_pptx_preview_html( *, deck_title: str, author: str | None, date_created: str | None, date_modified: str | None, slides: list[dict[str, object]], ) -> str: slide_sections = [] for slide in slides: slide_number = int(slide["slide_number"]) notes_blocks = list(slide.get("notes_blocks", [])) notes_section = "" if notes_blocks: notes_section = ( '

Speaker Notes

' f'{render_html_text_blocks([str(block) for block in notes_blocks])}
' ) slide_sections.append( f'
' f"

Slide {slide_number}

" f'{render_pptx_content_entries([dict(entry) for entry in slide["content_entries"]])}' f"{notes_section}" "
" ) metadata_rows = { passive_field_label("title"): deck_title, passive_field_label("author"): author or "", passive_field_label("date_created"): date_created or "", passive_field_label("date_modified"): date_modified or "", } return build_html_preview( metadata_rows, document_title=deck_title, head_html=( "" ), body_html=( "".join(slide_sections) ), ) def extract_pptx_file(path: Path) -> dict[str, object]: with zipfile.ZipFile(path) as archive: core_properties_root = None try: core_properties_root = parse_xml_document(archive.read("docProps/core.xml")) except KeyError: core_properties_root = None deck_title = None author = None subject = None date_created = None date_modified = None if core_properties_root is not None: deck_title = normalize_whitespace(core_properties_root.findtext("./dc:title", default="", namespaces=PPTX_NAMESPACES)) author = normalize_whitespace(core_properties_root.findtext("./dc:creator", default="", namespaces=PPTX_NAMESPACES)) or None subject = normalize_whitespace(core_properties_root.findtext("./dc:subject", default="", namespaces=PPTX_NAMESPACES)) or None date_created = normalize_datetime( core_properties_root.findtext("./dcterms:created", default="", namespaces=PPTX_NAMESPACES) ) date_modified = normalize_datetime( core_properties_root.findtext("./dcterms:modified", default="", namespaces=PPTX_NAMESPACES) ) presentation_root = parse_xml_document(archive.read("ppt/presentation.xml")) presentation_relationships = read_ooxml_relationships(archive, "ppt/presentation.xml") slide_part_names: list[str] = [] for slide_id in presentation_root.findall("./p:sldIdLst/p:sldId", PPTX_NAMESPACES): rel_id = slide_id.attrib.get(f"{{{PPTX_NAMESPACES['r']}}}id") if not rel_id: continue relationship = presentation_relationships.get(rel_id) if relationship is None: continue slide_part_names.append(relationship["target"]) slides: list[dict[str, object]] = [] text_sections: list[str] = [] for index, slide_part_name in enumerate(slide_part_names, start=1): slide_root = parse_xml_document(archive.read(slide_part_name)) slide_tree = slide_root.find("./p:cSld/p:spTree", PPTX_NAMESPACES) slide_relationships = read_ooxml_relationships(archive, slide_part_name) content_entries = ( sorted_pptx_content_entries(slide_tree, archive=archive, relationships=slide_relationships) if slide_tree is not None else [] ) text_blocks = [str(block) for entry in content_entries if entry.get("kind") == "text" for block in entry["blocks"]] notes_blocks = extract_pptx_notes_blocks(archive, slide_part_name) slides.append( { "slide_number": index, "content_entries": content_entries, "text_blocks": text_blocks, "notes_blocks": notes_blocks, } ) section_lines = [f"Slide {index}"] section_lines.extend(text_blocks) if notes_blocks: section_lines.append("Speaker notes") section_lines.extend(notes_blocks) text_sections.append("\n".join(line for line in section_lines if line)) if deck_title and deck_title.strip().lower() in {"powerpoint presentation", "presentation"}: deck_title = None resolved_title = deck_title or path.stem preview = build_pptx_preview_html( deck_title=resolved_title, author=author, date_created=date_created, date_modified=date_modified, slides=slides, ) text_content = normalize_whitespace("\n\n".join(section for section in text_sections if section)) return { "page_count": len(slides), "author": author, "content_type": "Presentation", "date_created": date_created, "date_modified": date_modified, "participants": None, "title": resolved_title, "subject": subject, "recipients": None, "text_content": text_content, "text_status": "empty" if not text_content else "ok", "preview_artifacts": [ { "file_name": f"{path.name}.html", "preview_type": "html", "label": "deck", "ordinal": 0, "content": preview, } ], } def slugify(value: str) -> str: slug = re.sub(r"[^a-zA-Z0-9]+", "-", value.strip().lower()).strip("-") return slug or "item" def preview_base_path_for_rel_path(rel_path: str) -> Path: source_rel_path = container_source_rel_path_from_message_rel_path(rel_path) if source_rel_path is not None: return Path("previews") / Path(source_rel_path) / "messages" base = Path(rel_path) if base.parts and base.parts[0] == INTERNAL_REL_PATH_PREFIX: base = Path(*base.parts[1:]) if base.parts and base.parts[0] == "previews": return base.parent return Path("previews") / base.parent def production_source_part_targets( paths: dict[str, Path], connection: sqlite3.Connection, row: sqlite3.Row | None, ) -> list[dict[str, object]]: if row is None or row["source_kind"] != PRODUCTION_SOURCE_KIND: return [] source_rows = connection.execute( """ SELECT part_kind, rel_source_path, ordinal, label FROM document_source_parts WHERE document_id = ? ORDER BY CASE part_kind WHEN 'native' THEN 0 WHEN 'image' THEN 1 ELSE 2 END, ordinal ASC, id ASC """, (row["id"],), ).fetchall() targets: list[dict[str, object]] = [] for source_row in source_rows: rel_source_path = str(source_row["rel_source_path"]) abs_path = paths["root"] / rel_source_path if not abs_path.exists(): continue part_kind = str(source_row["part_kind"]) preview_type = "native" if part_kind == "image": preview_type = "image" targets.append( { "rel_path": rel_source_path, "abs_path": str(abs_path), "preview_type": preview_type, "label": source_row["label"], "ordinal": int(source_row["ordinal"]), } ) return targets def document_native_target(paths: dict[str, Path], row: sqlite3.Row | None) -> dict[str, object] | None: if row is None: return None rel_path = str(row["rel_path"]) abs_path = document_absolute_path(paths, rel_path) if not abs_path.exists(): return None return { "rel_path": rel_path, "abs_path": str(abs_path), "preview_type": "native", "label": None, "ordinal": 0, } def document_prefers_native_primary_preview(row: sqlite3.Row | None) -> bool: if row is None: return False content_type = normalize_whitespace(str(row["content_type"] or "")) file_type = normalize_whitespace(str(row["file_type"] or "")).lower() if not file_type: file_type = normalize_extension(Path(str(row["file_name"] or row["rel_path"] or ""))) if content_type == "Chat": return file_type in {"pdf", "docx", "rtf"} return False def build_preview_target_payload( *, rel_path: str, abs_path: str, preview_type: str, label: str | None, ordinal: int, target_fragment: object = None, ) -> dict[str, object]: normalized_fragment = normalize_whitespace(str(target_fragment or "")) or None return { "rel_path": append_preview_fragment(rel_path, normalized_fragment), "abs_path": append_preview_fragment(abs_path, normalized_fragment), "file_rel_path": rel_path, "file_abs_path": abs_path, "preview_type": preview_type, "label": label, "ordinal": ordinal, "target_fragment": normalized_fragment, } def preview_target_payload_from_preview_row(paths: dict[str, Path], preview_row: sqlite3.Row) -> dict[str, object]: rel_preview = str(Path(INTERNAL_REL_PATH_PREFIX) / preview_row["rel_preview_path"]) abs_preview = str(paths["state_dir"] / preview_row["rel_preview_path"]) return build_preview_target_payload( rel_path=rel_preview, abs_path=abs_preview, preview_type=str(preview_row["preview_type"]), label=(str(preview_row["label"]) if preview_row["label"] is not None else None), ordinal=int(preview_row["ordinal"]), target_fragment=preview_row["target_fragment"] if "target_fragment" in preview_row.keys() else None, ) def preview_rows_use_conversation_navigation(preview_rows: list[sqlite3.Row]) -> bool: return bool(preview_rows) and is_conversation_preview_rel_path(preview_rows[0]["rel_preview_path"]) def ordered_preview_rows_for_document( document_row: sqlite3.Row | None, preview_rows: list[sqlite3.Row], ) -> list[sqlite3.Row]: if not preview_rows: return [] if document_row is None or not preview_rows_use_conversation_navigation(preview_rows): return list(preview_rows) if normalize_whitespace(str(document_row["content_type"] or "")) == "Chat": return list(preview_rows) for preview_row in preview_rows: if not is_conversation_preview_rel_path(preview_row["rel_preview_path"]): return [preview_row, *[row for row in preview_rows if row is not preview_row]] return list(preview_rows) def conversation_primary_preview_target( paths: dict[str, Path], document_row: sqlite3.Row | None, preview_rows: list[sqlite3.Row], ) -> dict[str, object] | None: if document_row is None or document_row["conversation_id"] is None: return None conversation_rows = [ preview_row for preview_row in preview_rows if is_conversation_preview_rel_path(str(preview_row["rel_preview_path"] or "")) and normalize_whitespace(str(preview_row["label"] or "")).lower() != "contents" ] if not conversation_rows: return None selected_row = next( ( preview_row for preview_row in conversation_rows if normalize_whitespace(str(preview_row["target_fragment"] or "")) ), conversation_rows[0], ) selected_target = preview_target_payload_from_preview_row(paths, selected_row) return selected_target def default_preview_target(paths: dict[str, Path], row: sqlite3.Row, connection: sqlite3.Connection) -> dict[str, object]: preview_rows = connection.execute( """ SELECT rel_preview_path, preview_type, target_fragment, label, ordinal FROM document_previews WHERE document_id = ? ORDER BY ordinal ASC, id ASC """, (row["id"],), ).fetchall() ordered_preview_rows = ordered_preview_rows_for_document(row, preview_rows) native_target = document_native_target(paths, row) if ( ordered_preview_rows and not preview_rows_use_conversation_navigation(ordered_preview_rows) and document_prefers_native_primary_preview(row) and native_target is not None ): return build_preview_target_payload( rel_path=str(native_target["rel_path"]), abs_path=str(native_target["abs_path"]), preview_type=str(native_target["preview_type"]), label=(str(native_target["label"]) if native_target["label"] is not None else None), ordinal=int(native_target["ordinal"]), ) conversation_target = conversation_primary_preview_target(paths, row, ordered_preview_rows) if conversation_target is not None: return conversation_target if ordered_preview_rows: return preview_target_payload_from_preview_row(paths, ordered_preview_rows[0]) source_targets = production_source_part_targets(paths, connection, row) if source_targets: return build_preview_target_payload( rel_path=str(source_targets[0]["rel_path"]), abs_path=str(source_targets[0]["abs_path"]), preview_type=str(source_targets[0]["preview_type"]), label=(str(source_targets[0]["label"]) if source_targets[0]["label"] is not None else None), ordinal=int(source_targets[0]["ordinal"]), ) rel_path = row["rel_path"] return build_preview_target_payload( rel_path=str(rel_path), abs_path=str(document_absolute_path(paths, rel_path)), preview_type="native", label=None, ordinal=0, ) def collect_preview_targets(paths: dict[str, Path], document_id: int, rel_path: str, connection: sqlite3.Connection) -> list[dict[str, object]]: document_row = connection.execute("SELECT * FROM documents WHERE id = ?", (document_id,)).fetchone() preview_rows = connection.execute( """ SELECT rel_preview_path, preview_type, target_fragment, label, ordinal FROM document_previews WHERE document_id = ? ORDER BY ordinal ASC, id ASC """, (document_id,), ).fetchall() if not preview_rows: native_target = document_native_target(paths, document_row) source_targets = production_source_part_targets(paths, connection, document_row) if native_target is not None and document_prefers_native_primary_preview(document_row): targets = [native_target] for target in source_targets: if target["rel_path"] not in {existing["rel_path"] for existing in targets}: targets.append(target) return targets if source_targets: return source_targets if native_target is not None: return [native_target] abs_path = document_absolute_path(paths, rel_path) return [ build_preview_target_payload( rel_path=rel_path, abs_path=str(abs_path), preview_type="native", label=None, ordinal=0, ) ] ordered_preview_rows = ordered_preview_rows_for_document(document_row, preview_rows) targets: list[dict[str, object]] = [] seen_targets: set[tuple[str, str | None]] = set() def append_target(target: dict[str, object]) -> None: key = ( str(target["file_rel_path"]), normalize_whitespace(str(target.get("target_fragment") or "")) or None, ) if key in seen_targets: return seen_targets.add(key) targets.append(target) if ( ordered_preview_rows and not preview_rows_use_conversation_navigation(ordered_preview_rows) and document_prefers_native_primary_preview(document_row) ): native_target = document_native_target(paths, document_row) if native_target is not None: append_target( build_preview_target_payload( rel_path=str(native_target["rel_path"]), abs_path=str(native_target["abs_path"]), preview_type=str(native_target["preview_type"]), label=(str(native_target["label"]) if native_target["label"] is not None else None), ordinal=int(native_target["ordinal"]), ) ) conversation_target = conversation_primary_preview_target(paths, document_row, ordered_preview_rows) if conversation_target is not None: append_target(conversation_target) for preview_row in ordered_preview_rows: preview_target = preview_target_payload_from_preview_row(paths, preview_row) if ( conversation_target is not None and str(preview_target["file_rel_path"]) == str(conversation_target["file_rel_path"]) and normalize_whitespace(str(preview_target.get("label") or "")).lower() in {"conversation", "segment"} ): continue append_target(preview_target) source_targets = production_source_part_targets(paths, connection, document_row) for target in source_targets: if target["rel_path"] not in {existing["rel_path"] for existing in targets}: append_target( build_preview_target_payload( rel_path=str(target["rel_path"]), abs_path=str(target["abs_path"]), preview_type=str(target["preview_type"]), label=(str(target["label"]) if target["label"] is not None else None), ordinal=int(target["ordinal"]), ) ) return targets def token_estimate(text: str) -> int: return max(1, len(text) // 4) def chunk_text(text: str, max_chars: int = CHUNK_TARGET_CHARS, overlap: int = CHUNK_OVERLAP_CHARS) -> list[dict[str, object]]: normalized = normalize_whitespace(text) if not normalized: return [] chunks: list[dict[str, object]] = [] start = 0 length = len(normalized) chunk_index = 0 while start < length: end = min(length, start + max_chars) if end < length: preferred_break = normalized.rfind("\n", max(start + 400, end - 500), end) if preferred_break <= start: preferred_break = normalized.rfind(" ", max(start + 400, end - 250), end) if preferred_break > start: end = preferred_break chunk_text_value = normalized[start:end].strip() if chunk_text_value: chunk_start = normalized.find(chunk_text_value, start, end + len(chunk_text_value)) chunk_end = chunk_start + len(chunk_text_value) chunks.append( { "chunk_index": chunk_index, "char_start": chunk_start, "char_end": chunk_end, "token_estimate": token_estimate(chunk_text_value), "text_content": chunk_text_value, } ) chunk_index += 1 if end >= length: break start = max(end - overlap, start + 1) return chunks def extracted_search_chunks(extracted: dict[str, object]) -> list[dict[str, object]]: raw_chunks = extracted.get("chunks") if isinstance(raw_chunks, list): normalized_chunks: list[dict[str, object]] = [] for index, raw_chunk in enumerate(raw_chunks): if not isinstance(raw_chunk, dict): continue text_content = normalize_whitespace(str(raw_chunk.get("text_content") or "")) if not text_content: continue normalized_chunks.append( { "chunk_index": int(raw_chunk.get("chunk_index", index)), "char_start": int(raw_chunk.get("char_start", 0)), "char_end": int(raw_chunk.get("char_end", len(text_content))), "token_estimate": int(raw_chunk.get("token_estimate", token_estimate(text_content))), "text_content": text_content, } ) return normalized_chunks return chunk_text(str(extracted.get("text_content") or "")) SPREADSHEET_TYPE_SAMPLE_LIMIT = 50 SPREADSHEET_MAX_COLUMNS_PER_SHEET = 100 SPREADSHEET_MAX_COMMENTS_PER_SHEET = 200 SPREADSHEET_MAX_HYPERLINKS_PER_WORKBOOK = 200 SPREADSHEET_MAX_NAMED_RANGES = 500 SPREADSHEET_MAX_ENUM_VALUES = 64 SPREADSHEET_MAX_PARTICIPANTS = 32 SPREADSHEET_MAX_SUMMARY_CHARS = 64 * 1024 SPREADSHEET_HTML_PREVIEW_MAX_SOURCE_BYTES = 10 * 1024 * 1024 SPREADSHEET_HTML_PREVIEW_SHEETJS_URL = "https://cdnjs.cloudflare.com/ajax/libs/xlsx/0.18.5/xlsx.full.min.js" SPREADSHEET_XLSX_READ_ONLY_FALLBACK_BYTES = 50 * 1024 * 1024 SLACK_USER_DIRECTORY_CACHE: dict[str, dict[str, dict[str, str | None]]] = {} SLACK_USER_MENTION_PATTERN = re.compile(r"<@([A-Z0-9]+)(?:\|[^>]+)?>") SLACK_CHANNEL_MENTION_PATTERN = re.compile(r"<#([A-Z0-9]+)(?:\|([^>]+))?>") SLACK_SPECIAL_MENTION_PATTERN = re.compile(r"]+)?>") SLACK_NAMED_LINK_PATTERN = re.compile(r"<([^>|]+)\|([^>]+)>") SLACK_BARE_LINK_PATTERN = re.compile(r"<(https?://[^>]+)>") SLACK_EMOJI_SHORTCODE_PATTERN = re.compile(r":([a-z0-9_+\-]+):") SLACK_EMOJI_NAME_ALIASES = { "+1": "THUMBS UP SIGN", "-1": "THUMBS DOWN SIGN", "100": "HUNDRED POINTS SYMBOL", "christmas_tree": "🎄", "clap": "CLAPPING HANDS SIGN", "eyes": "EYES", "fire": "FIRE", "grin": "GRINNING FACE WITH SMILING EYES", "heart": "HEAVY BLACK HEART", "joy": "FACE WITH TEARS OF JOY", "ok_hand": "OK HAND SIGN", "partying_face": "🥳", "pray": "PERSON WITH FOLDED HANDS", "rocket": "ROCKET", "smile": "SMILING FACE WITH OPEN MOUTH", "sob": "LOUDLY CRYING FACE", "tada": "PARTY POPPER", "thinking_face": "THINKING FACE", "thumbsdown": "THUMBS DOWN SIGN", "thumbsup": "THUMBS UP SIGN", "warning": "WARNING SIGN", "wave": "WAVING HAND SIGN", "white_check_mark": "WHITE HEAVY CHECK MARK", } def choose_slack_text(*values: object) -> str | None: for value in values: normalized = normalize_whitespace(str(value or "")) if normalized: return normalized return None def replace_slack_emoji_shortcodes(text: str) -> str: def replace_match(match: re.Match[str]) -> str: token = match.group(1).lower() candidate_names = [ SLACK_EMOJI_NAME_ALIASES.get(token), token.replace("_", " ").replace("-", " ").upper(), ] for candidate_name in candidate_names: if not candidate_name: continue if any(ord(character) > 127 for character in candidate_name): return candidate_name try: return unicodedata.lookup(candidate_name) except KeyError: continue return match.group(0) return SLACK_EMOJI_SHORTCODE_PATTERN.sub(replace_match, text) def normalize_slack_user_info( user_record: dict[str, object] | None = None, inline_profile: dict[str, object] | None = None, ) -> dict[str, str | None]: record = user_record or {} profile = dict(record.get("profile") if isinstance(record.get("profile"), dict) else {}) if inline_profile: for key, value in inline_profile.items(): if value not in (None, ""): profile[key] = value first_name = choose_slack_text(profile.get("first_name"), record.get("first_name")) last_name = choose_slack_text(profile.get("last_name"), record.get("last_name")) combined_name = " ".join(part for part in [first_name, last_name] if part) real_name = choose_slack_text(profile.get("real_name"), record.get("real_name")) display_name = choose_slack_text(profile.get("display_name")) handle = choose_slack_text(profile.get("name"), record.get("name")) email = normalize_entity_email(choose_slack_text(profile.get("email"), record.get("email")) or "") speaker_name = choose_slack_text( real_name, display_name, combined_name, handle, ) mention_name = choose_slack_text( display_name, first_name, handle, speaker_name, ) return { "avatar_color": choose_slack_text(profile.get("color"), record.get("color")), "email": email, "handle": handle, "real_name": real_name, "display_name": display_name, "speaker_name": speaker_name, "mention_name": mention_name, } def slack_export_root_for_path(path: Path) -> Path | None: current = path.parent while True: if (current / "users.json").exists(): return current if current.parent == current: return None current = current.parent def load_slack_user_directory(export_root: Path | None) -> dict[str, dict[str, str | None]]: if export_root is None: return {} cache_key = str(export_root) cached = SLACK_USER_DIRECTORY_CACHE.get(cache_key) if cached is not None: return cached users_path = export_root / "users.json" directory: dict[str, dict[str, str | None]] = {} try: raw_users = json.loads(users_path.read_text(encoding="utf-8")) except (OSError, ValueError, json.JSONDecodeError): raw_users = None if isinstance(raw_users, list): for item in raw_users: if not isinstance(item, dict): continue user_id = choose_slack_text(item.get("id")) if not user_id: continue directory[user_id] = normalize_slack_user_info(item) SLACK_USER_DIRECTORY_CACHE[cache_key] = directory return directory def resolve_slack_user_info( user_id: str | None, user_directory: dict[str, dict[str, str | None]], inline_profile: dict[str, object] | None = None, ) -> dict[str, str | None]: resolved = dict(user_directory.get(user_id or "", {})) inline_info = normalize_slack_user_info({}, inline_profile) if inline_profile else {} for key, value in inline_info.items(): if value: resolved[key] = value if user_id and not resolved.get("speaker_name"): resolved["speaker_name"] = user_id if user_id and not resolved.get("mention_name"): resolved["mention_name"] = user_id if user_id: resolved["slack_user_id"] = user_id return resolved def format_slack_document_title(path: Path) -> str: channel_name = normalize_whitespace(path.parent.name) if channel_name and not channel_name.startswith("#"): channel_name = f"#{channel_name}" day_token = normalize_whitespace(path.stem) try: day_label = date.fromisoformat(day_token).strftime("%b %d, %Y").replace(" 0", " ") except ValueError: day_label = day_token if channel_name and day_label: return f"{channel_name} - {day_label}" return channel_name or day_label or "Slack conversation" def normalize_slack_timestamp(value: object) -> str | None: raw = normalize_whitespace(str(value or "")) if not raw: return None try: seconds = float(raw) except (TypeError, ValueError): return None return datetime.fromtimestamp(seconds, timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z") def render_slack_text(text: str, user_directory: dict[str, dict[str, str | None]]) -> str: rendered = str(text or "") rendered = SLACK_USER_MENTION_PATTERN.sub( lambda match: f"@{resolve_slack_user_info(match.group(1), user_directory).get('mention_name') or match.group(1)}", rendered, ) rendered = SLACK_CHANNEL_MENTION_PATTERN.sub(lambda match: f"#{match.group(2) or match.group(1)}", rendered) rendered = SLACK_SPECIAL_MENTION_PATTERN.sub(lambda match: f"@{match.group(1)}", rendered) rendered = SLACK_NAMED_LINK_PATTERN.sub(lambda match: match.group(2), rendered) rendered = SLACK_BARE_LINK_PATTERN.sub(lambda match: match.group(1), rendered) return normalize_whitespace(replace_slack_emoji_shortcodes(rendered)) def slack_message_actor_info( message: dict[str, object], user_directory: dict[str, dict[str, str | None]], ) -> dict[str, str | None]: user_id = choose_slack_text(message.get("user")) user_profile = message.get("user_profile") if isinstance(message.get("user_profile"), dict) else None if user_id: return resolve_slack_user_info(user_id, user_directory, user_profile) bot_profile = message.get("bot_profile") if isinstance(message.get("bot_profile"), dict) else None if bot_profile: actor_info = normalize_slack_user_info({}, bot_profile) if not actor_info.get("speaker_name"): actor_info["speaker_name"] = "Slack bot" return actor_info speaker_name = ( choose_slack_text(message.get("username")) or choose_slack_text(message.get("subtype")) or "Slack message" ) return { "avatar_color": None, "mention_name": speaker_name, "speaker_name": speaker_name, } def iter_slack_export_entries( raw_value: object, user_directory: dict[str, dict[str, str | None]], ) -> list[dict[str, object]]: candidate_items = raw_value.get("messages") if isinstance(raw_value, dict) else raw_value if not isinstance(candidate_items, list): return [] entries: list[dict[str, object]] = [] for item in candidate_items: if not isinstance(item, dict): continue if normalize_whitespace(str(item.get("type") or "")).lower() != "message": continue body = render_slack_text(choose_slack_text(item.get("text")) or "", user_directory) if not body: continue actor_info = slack_message_actor_info(item, user_directory) speaker = actor_info.get("speaker_name") or "Slack message" timestamp = normalize_slack_timestamp(item.get("ts")) entries.append( { "avatar_color": actor_info.get("avatar_color"), "speaker": speaker, "body": body, "timestamp": timestamp, "timestamp_label": format_chat_preview_timestamp(timestamp), "avatar_label": chat_avatar_initials(speaker), } ) return entries def extract_slack_chat_json_payload(path: Path, decoded_text: str) -> dict[str, object] | None: try: raw_value = json.loads(decoded_text) except (TypeError, ValueError, json.JSONDecodeError): return None candidate_items = raw_value.get("messages") if isinstance(raw_value, dict) else raw_value if not isinstance(candidate_items, list): return None message_like_count = sum( 1 for item in candidate_items if isinstance(item, dict) and normalize_whitespace(str(item.get("type") or "")).lower() == "message" and item.get("ts") is not None and any(item.get(key) is not None for key in ("text", "user", "user_profile", "subtype")) ) if message_like_count == 0: return None user_directory = load_slack_user_directory(slack_export_root_for_path(path)) entries = iter_slack_export_entries(raw_value, user_directory) if not entries: return None participants: list[str] = [] seen: set[str] = set() first_body: str | None = None first_timestamp: str | None = None last_timestamp: str | None = None timestamped_matches = 0 transcript_lines: list[str] = [] for entry in entries: speaker = normalize_whitespace(str(entry.get("speaker") or "")) or "Unknown" key = speaker.lower() if key not in seen: seen.add(key) participants.append(speaker) body = str(entry.get("body") or "").strip() if first_body is None and body: first_body = body timestamp = entry.get("timestamp") if isinstance(timestamp, str) and timestamp: timestamped_matches += 1 if first_timestamp is None: first_timestamp = timestamp last_timestamp = timestamp transcript_lines.append(f"[{timestamp}] {speaker}: {body}") else: transcript_lines.append(f"{speaker}: {body}") return { "text_content": normalize_whitespace("\n".join(transcript_lines)), "chat_metadata": { "author": None, "participants": ", ".join(participants) or None, "date_created": first_timestamp, "date_modified": last_timestamp if last_timestamp and last_timestamp != first_timestamp else None, "title": format_slack_document_title(path), "message_count": len(entries), "timestamped_message_count": timestamped_matches, }, "chat_entries": entries, } def extract_plain_text_file(path: Path) -> dict[str, object]: decoded, text_status, _ = decode_bytes(path.read_bytes()) file_type = normalize_extension(path) structured_chat = extract_slack_chat_json_payload(path, decoded) if file_type == "json" else None text_content = ( str(structured_chat["text_content"]) if structured_chat and structured_chat.get("text_content") else (strip_html_tags(decoded) if file_type in {"htm", "html"} else normalize_whitespace(decoded)) ) email_headers = {} if structured_chat else extract_email_like_headers(text_content) chat_metadata = ( dict(structured_chat["chat_metadata"]) if structured_chat and isinstance(structured_chat.get("chat_metadata"), dict) else extract_chat_transcript_metadata(text_content) ) chat_entries = structured_chat.get("chat_entries") if structured_chat and isinstance(structured_chat.get("chat_entries"), list) else None participants = extract_email_chain_participants( text_content, [email_headers.get("author"), email_headers.get("recipients")] if email_headers else None, ) or (str(chat_metadata["participants"]) if chat_metadata and chat_metadata.get("participants") else None) or extract_chat_participants(text_content) title = email_headers.get("title") if email_headers else (str(chat_metadata["title"]) if chat_metadata and chat_metadata.get("title") else None) if title is None and file_type in {"md", "txt"} and text_content: title = text_content.splitlines()[0][:200] resolved_author = None if chat_metadata else (email_headers.get("author") if email_headers else None) preview_artifacts = ( build_chat_preview_artifacts( title=title, author=resolved_author, participants=participants, date_created=(email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None), date_modified=str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None, text_body=text_content, preview_file_name=f"{path.name}.html", chat_metadata=chat_metadata, chat_entries=chat_entries, ) if chat_metadata else [] ) return { "page_count": None, "author": resolved_author, "content_type": determine_content_type(path, text_content, email_headers=email_headers, chat_metadata=chat_metadata), "date_created": (email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None), "date_modified": str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None, "participants": participants, "title": title, "subject": email_headers.get("subject") if email_headers else None, "recipients": email_headers.get("recipients") if email_headers else None, "text_content": text_content, "text_status": "empty" if not text_content else text_status, "preview_artifacts": preview_artifacts, } def extract_native_preview_only_file(path: Path, explicit_content_type: str | None = None) -> dict[str, object]: return { "page_count": None, "author": None, "content_type": explicit_content_type or determine_content_type(path, ""), "date_created": None, "date_modified": None, "participants": None, "title": None, "subject": None, "recipients": None, "text_content": "", "text_status": "empty", "preview_artifacts": [], } def extract_rtf_file(path: Path) -> dict[str, object]: rtf_to_text_fn = dependency_guard("rtf_to_text", "striprtf", "rtf") decoded, text_status, _ = decode_bytes(path.read_bytes()) text_content = normalize_whitespace(rtf_to_text_fn(decoded)) email_headers = extract_email_like_headers(text_content) chat_metadata = extract_chat_transcript_metadata(text_content) participants = extract_email_chain_participants( text_content, [email_headers.get("author"), email_headers.get("recipients")] if email_headers else None, ) or (str(chat_metadata["participants"]) if chat_metadata and chat_metadata.get("participants") else None) or extract_chat_participants(text_content) title = email_headers.get("title") if email_headers else (str(chat_metadata["title"]) if chat_metadata and chat_metadata.get("title") else None) if title is None and text_content: title = text_content.splitlines()[0][:200] resolved_author = None if chat_metadata else (email_headers.get("author") if email_headers else None) preview_artifacts = ( build_chat_preview_artifacts( title=title, author=resolved_author, participants=participants, date_created=(email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None), date_modified=str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None, text_body=text_content, preview_file_name=f"{path.name}.html", chat_metadata=chat_metadata, label="text", ) if chat_metadata else [ { "file_name": f"{path.name}.html", "preview_type": "html", "label": "text", "ordinal": 0, "content": build_html_preview( {}, body_text=text_content, document_title=title or path.stem or path.name, ), } ] ) return { "page_count": None, "author": resolved_author, "content_type": determine_content_type( path, text_content, email_headers=email_headers, chat_metadata=chat_metadata, explicit_content_type="E-Doc", ), "date_created": (email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None), "date_modified": str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None, "participants": participants, "title": title, "subject": email_headers.get("subject") if email_headers else None, "recipients": email_headers.get("recipients") if email_headers else None, "text_content": text_content, "text_status": "empty" if not text_content else text_status, "preview_artifacts": preview_artifacts, } def extract_pdf_file(path: Path) -> dict[str, object]: pdfplumber_module = dependency_guard("pdfplumber", "pdfplumber", "pdf") with pdfplumber_module.open(path) as pdf: # type: ignore[union-attr] metadata = pdf.metadata or {} texts = [(page.extract_text() or "").strip() for page in pdf.pages] text_content = normalize_whitespace("\n\n".join(part for part in texts if part)) email_headers = extract_email_like_headers(texts[0] if texts else "") chat_metadata = extract_chat_transcript_metadata(text_content) participants = extract_email_chain_participants( text_content, [email_headers.get("author"), email_headers.get("recipients")] if email_headers else None, ) or (str(chat_metadata["participants"]) if chat_metadata and chat_metadata.get("participants") else None) or extract_chat_participants(text_content) resolved_author = ( None if chat_metadata else (email_headers.get("author") if email_headers else None) or metadata.get("Author") ) preview_artifacts = ( build_chat_preview_artifacts( title=(email_headers.get("title") if email_headers else None) or (str(chat_metadata["title"]) if chat_metadata and chat_metadata.get("title") else None) or metadata.get("Title"), author=resolved_author, participants=participants, date_created=(email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None) or normalize_datetime(metadata.get("CreationDate")), date_modified=(str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None) or normalize_datetime(metadata.get("ModDate")), text_body=text_content, preview_file_name=f"{path.name}.html", chat_metadata=chat_metadata, ) if chat_metadata else [] ) return { "page_count": len(pdf.pages), "author": resolved_author, "content_type": determine_content_type( path, text_content, email_headers=email_headers, chat_metadata=chat_metadata, explicit_content_type="E-Doc", ), "date_created": (email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None) or normalize_datetime(metadata.get("CreationDate")), "date_modified": (str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None) or normalize_datetime(metadata.get("ModDate")), "participants": participants, "title": (email_headers.get("title") if email_headers else None) or (str(chat_metadata["title"]) if chat_metadata and chat_metadata.get("title") else None) or metadata.get("Title"), "subject": (email_headers.get("subject") if email_headers else None) or metadata.get("Subject"), "recipients": email_headers.get("recipients") if email_headers else None, "text_content": text_content, "text_status": "empty" if not text_content else "ok", "preview_artifacts": preview_artifacts, } def extract_docx_file(path: Path) -> dict[str, object]: docx_document_cls = dependency_guard("DocxDocument", "python-docx", "docx") document = docx_document_cls(path) # type: ignore[operator] text_content = normalize_whitespace("\n\n".join(paragraph.text for paragraph in document.paragraphs if paragraph.text)) props = document.core_properties email_headers = extract_email_like_headers(text_content) chat_metadata = extract_chat_transcript_metadata(text_content) participants = extract_email_chain_participants( text_content, [email_headers.get("author"), email_headers.get("recipients")] if email_headers else None, ) or (str(chat_metadata["participants"]) if chat_metadata and chat_metadata.get("participants") else None) or extract_chat_participants(text_content) resolved_author = ( None if chat_metadata else (email_headers.get("author") if email_headers else None) or props.author or None ) preview_artifacts = ( build_chat_preview_artifacts( title=(email_headers.get("title") if email_headers else None) or (str(chat_metadata["title"]) if chat_metadata and chat_metadata.get("title") else None) or props.title or None, author=resolved_author, participants=participants, date_created=(email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None) or normalize_datetime(props.created), date_modified=(str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None) or normalize_datetime(props.modified), text_body=text_content, preview_file_name=f"{path.name}.html", chat_metadata=chat_metadata, ) if chat_metadata else [] ) return { "page_count": None, "author": resolved_author, "content_type": determine_content_type( path, text_content, email_headers=email_headers, chat_metadata=chat_metadata, explicit_content_type="E-Doc", ), "date_created": (email_headers.get("date_created") if email_headers else None) or (str(chat_metadata["date_created"]) if chat_metadata and chat_metadata.get("date_created") else None) or normalize_datetime(props.created), "date_modified": (str(chat_metadata["date_modified"]) if chat_metadata and chat_metadata.get("date_modified") else None) or normalize_datetime(props.modified), "participants": participants, "title": (email_headers.get("title") if email_headers else None) or (str(chat_metadata["title"]) if chat_metadata and chat_metadata.get("title") else None) or props.title or None, "subject": (email_headers.get("subject") if email_headers else None) or props.subject or None, "recipients": email_headers.get("recipients") if email_headers else None, "text_content": text_content, "text_status": "empty" if not text_content else "ok", "preview_artifacts": preview_artifacts, } def stringify_spreadsheet_value(value: object) -> object: if value is None: return "" if isinstance(value, bool): return "TRUE" if value else "FALSE" if isinstance(value, datetime): return normalize_datetime(value) or value.isoformat() if isinstance(value, date): return value.isoformat() if isinstance(value, float) and value.is_integer(): return str(int(value)) return value def spreadsheet_column_label(column_index: int) -> str: index = max(1, int(column_index)) letters: list[str] = [] while index > 0: index, remainder = divmod(index - 1, 26) letters.append(chr(ord("A") + remainder)) return "".join(reversed(letters)) def spreadsheet_text_value(value: object) -> str: rendered = stringify_spreadsheet_value(value) if rendered == "": return "" return normalize_inline_whitespace(str(rendered)) def spreadsheet_value_is_present(value: object) -> bool: if value is None: return False if isinstance(value, str): return bool(value.strip()) return True def spreadsheet_row_last_nonempty_index(values: list[object]) -> int: for index in range(len(values), 0, -1): if spreadsheet_value_is_present(values[index - 1]): return index return 0 def normalize_spreadsheet_headers(raw_headers: list[object], column_count: int) -> list[str]: headers: list[str] = [] seen: dict[str, int] = {} for column_index in range(1, column_count + 1): raw_value = raw_headers[column_index - 1] if column_index - 1 < len(raw_headers) else None base = spreadsheet_text_value(raw_value) or f"Column {spreadsheet_column_label(column_index)}" key = base.casefold() seen[key] = seen.get(key, 0) + 1 headers.append(base if seen[key] == 1 else f"{base} ({seen[key]})") return headers def append_spreadsheet_participants( participants: list[str], seen: set[str], raw_values: list[str | None], ) -> None: if len(participants) >= SPREADSHEET_MAX_PARTICIPANTS: return append_unique_participants(participants, seen, raw_values) if len(participants) > SPREADSHEET_MAX_PARTICIPANTS: del participants[SPREADSHEET_MAX_PARTICIPANTS:] def spreadsheet_number_format_hint(number_formats: list[str]) -> str | None: for raw_format in number_formats: normalized = str(raw_format or "").strip().lower() if not normalized or normalized == "general": continue if "%" in normalized: return "percent" if "[$" in normalized or any(symbol in normalized for symbol in ("$", "€", "£", "¥", "₹")): return "currency" if any(token in normalized for token in ("yy", "dd", "mmm", "am/pm")): if any(token in normalized for token in ("h", "ss", "am/pm")): return "datetime" return "date" return None def spreadsheet_string_shape(text: str) -> str: normalized = text.strip() lowered = normalized.lower() if not normalized: return "string" if lowered in {"true", "false", "yes", "no", "y", "n"}: return "boolean" if re.fullmatch(r"[$€£¥₹]\s*[-+]?\d[\d,]*(?:\.\d+)?", normalized): return "currency" if re.fullmatch(r"[-+]?\d+(?:\.\d+)?%", normalized): return "percent" if re.fullmatch(r"[-+]?\d+", normalized): return "integer" if re.fullmatch(r"[-+]?(?:\d+\.\d*|\d*\.\d+)(?:[eE][-+]?\d+)?", normalized) or re.fullmatch( r"[-+]?\d+[eE][-+]?\d+", normalized, ): return "number" if re.fullmatch(r"\d{4}-\d{2}-\d{2}[ T]\d{1,2}:\d{2}(?::\d{2})?(?:Z|[+-]\d{2}:\d{2})?", normalized): return "datetime" if re.fullmatch(r"\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2}(?::\d{2})?\s*(?:AM|PM)?", normalized, re.IGNORECASE): return "datetime" if re.fullmatch(r"\d{4}-\d{2}-\d{2}", normalized): return "date" if re.fullmatch(r"\d{1,2}/\d{1,2}/\d{2,4}", normalized): return "date" return "string" def infer_spreadsheet_value_kind(value: object) -> str: if isinstance(value, bool): return "boolean" if isinstance(value, datetime): return "datetime" if isinstance(value, date): return "date" if isinstance(value, int): return "integer" if isinstance(value, float): return "integer" if value.is_integer() else "number" if isinstance(value, str): return spreadsheet_string_shape(value) return "string" def infer_spreadsheet_column_type( samples: list[object], number_formats: list[str], *, force_enum: bool = False, ) -> str: if force_enum: return "enum" format_hint = spreadsheet_number_format_hint(number_formats) if format_hint is not None: return format_hint sample_kinds = {infer_spreadsheet_value_kind(value) for value in samples if spreadsheet_value_is_present(value)} if not sample_kinds: return "string" if sample_kinds <= {"boolean"}: return "boolean" if sample_kinds <= {"integer"}: return "integer" if sample_kinds <= {"integer", "number"}: return "number" if "number" in sample_kinds else "integer" if sample_kinds <= {"date"}: return "date" if sample_kinds <= {"date", "datetime"}: return "datetime" if "datetime" in sample_kinds else "date" if "currency" in sample_kinds and sample_kinds <= {"currency", "integer", "number"}: return "currency" if "percent" in sample_kinds and sample_kinds <= {"percent", "integer", "number"}: return "percent" return "string" def build_spreadsheet_sheet_section( sheet_name: str, column_descriptors: list[str], *, hidden_column_count: int = 0, comment_lines: list[str] | None = None, comment_overflow: int = 0, validation_lines: list[str] | None = None, hyperlink_lines: list[str] | None = None, chart_lines: list[str] | None = None, ) -> str: lines = [f"Sheet: {sheet_name}"] if column_descriptors: lines.append(f"Columns: {', '.join(column_descriptors)}") else: lines.append("Columns: none detected") if hidden_column_count > 0: lines.append(f"Columns truncated: {hidden_column_count} more columns not listed.") if comment_lines: lines.append("Comments:") lines.extend(comment_lines) if comment_overflow > 0: lines.append(f"Comments truncated: {comment_overflow} more comments not listed.") if validation_lines: lines.append("Validations:") lines.extend(validation_lines) if hyperlink_lines: lines.append("Hyperlinks:") lines.extend(hyperlink_lines) if chart_lines: lines.append("Charts:") lines.extend(chart_lines) return normalize_whitespace("\n".join(lines)) def build_spreadsheet_header_section( *, title: str, sheet_names: list[str], author: str | None, subject: str | None, participants: str | None, parse_note: str | None = None, ) -> str: lines = [ f"Workbook: {title}", f"Sheets: {', '.join(sheet_names) if sheet_names else title}", ] if author: lines.append(f"Author: {author}") if subject: lines.append(f"Subject: {subject}") if participants: lines.append(f"Participants: {participants}") if parse_note: lines.append(parse_note) return normalize_whitespace("\n".join(lines)) def build_structural_summary_from_sections(sections: list[str]) -> tuple[str, list[dict[str, object]]]: normalized_sections = [normalize_whitespace(section) for section in sections if normalize_whitespace(section)] if not normalized_sections: return "", [] kept_sections: list[str] = [] truncation_note = "Structural summary truncated." for section in normalized_sections: candidate_text = normalize_whitespace("\n\n".join([*kept_sections, section])) if len(candidate_text) <= SPREADSHEET_MAX_SUMMARY_CHARS: kept_sections.append(section) continue candidate_with_note = normalize_whitespace("\n\n".join([*kept_sections, truncation_note])) while kept_sections and len(candidate_with_note) > SPREADSHEET_MAX_SUMMARY_CHARS: kept_sections.pop() candidate_with_note = normalize_whitespace("\n\n".join([*kept_sections, truncation_note])) if len(candidate_with_note) <= SPREADSHEET_MAX_SUMMARY_CHARS: kept_sections.append(truncation_note) else: kept_sections = [truncation_note[:SPREADSHEET_MAX_SUMMARY_CHARS]] break text_content = normalize_whitespace("\n\n".join(kept_sections)) chunks: list[dict[str, object]] = [] chunk_index = 0 char_cursor = 0 for section_index, section in enumerate(kept_sections): if section_index > 0: char_cursor += 2 for section_chunk in chunk_text(section): chunks.append( { "chunk_index": chunk_index, "char_start": char_cursor + int(section_chunk["char_start"]), "char_end": char_cursor + int(section_chunk["char_end"]), "token_estimate": int(section_chunk["token_estimate"]), "text_content": str(section_chunk["text_content"]), } ) chunk_index += 1 char_cursor += len(section) return text_content, chunks def format_spreadsheet_preview_size(byte_count: int) -> str: size = float(max(0, int(byte_count))) units = ("bytes", "KB", "MB", "GB", "TB") unit_index = 0 while size >= 1024.0 and unit_index < len(units) - 1: size /= 1024.0 unit_index += 1 if unit_index == 0: return f"{int(size)} bytes" return f"{size:.1f} {units[unit_index]}" def build_spreadsheet_embedded_preview_html( *, document_title: str, file_name: str, file_type: str, source_base64: str | None, skip_reason: str | None = None, ) -> str: resolved_title = normalize_generated_document_title(document_title) or file_name or "Spreadsheet Preview" resolved_file_name = normalize_whitespace(file_name) or resolved_title resolved_file_type = normalize_whitespace(file_type).lower() resolved_skip_reason = normalize_whitespace(str(skip_reason or "")) or None return "".join( [ "" "", f"{html.escape(resolved_title)}", "", f'', "", "
" "
" "
" "
" "

Spreadsheet preview

" "

" "
" "

" "
" "
" "

Preparing preview...

" "
" "
" "" "
" "
" "
", ] ) def build_spreadsheet_html_preview_artifact(path: Path, *, ordinal: int = 0) -> dict[str, object]: file_size = path.stat().st_size source_base64: str | None = None skip_reason: str | None = None if file_size > SPREADSHEET_HTML_PREVIEW_MAX_SOURCE_BYTES: skip_reason = ( f"{path.name} is {format_spreadsheet_preview_size(file_size)}, which exceeds the " f"{format_spreadsheet_preview_size(SPREADSHEET_HTML_PREVIEW_MAX_SOURCE_BYTES)} embedded preview cap." ) else: source_base64 = base64.b64encode(path.read_bytes()).decode("ascii") return { "file_name": f"{path.name}.html", "preview_type": "html", "label": "spreadsheet", "ordinal": ordinal, "content": build_spreadsheet_embedded_preview_html( document_title=path.stem or path.name, file_name=path.name, file_type=normalize_extension(path) or "", source_base64=source_base64, skip_reason=skip_reason, ), } def parse_spreadsheet_literal_list(formula: str) -> list[str]: inner = formula[1:-1] try: values = next(csv.reader([inner])) except Exception: values = [part.strip() for part in inner.split(",")] deduped: list[str] = [] seen: set[str] = set() for value in values: normalized = spreadsheet_text_value(value) if not normalized: continue key = normalized.casefold() if key not in seen: seen.add(key) deduped.append(normalized) if len(deduped) >= SPREADSHEET_MAX_ENUM_VALUES: break return deduped def flatten_openpyxl_range_cells(range_cells: object) -> list[object]: if range_cells is None: return [] if isinstance(range_cells, tuple): flattened: list[object] = [] for item in range_cells: flattened.extend(flatten_openpyxl_range_cells(item)) return flattened return [range_cells] def resolve_openpyxl_reference_values(workbook: object, reference: object) -> list[str]: normalized_reference = normalize_inline_whitespace(str(reference or "")).lstrip("=") if not normalized_reference: return [] values: list[str] = [] if normalized_reference.startswith('"') and normalized_reference.endswith('"'): return parse_spreadsheet_literal_list(normalized_reference) destinations: list[tuple[str, str]] = [] direct_match = re.fullmatch(r"(?:'((?:[^']|'')+)'|([^!]+))!(.+)", normalized_reference) if direct_match: sheet_name = direct_match.group(1) or direct_match.group(2) or "" destinations.append((sheet_name.replace("''", "'"), direct_match.group(3))) else: defined_name = getattr(workbook, "defined_names", {}).get(normalized_reference) if defined_name is not None: try: destinations.extend(list(defined_name.destinations)) except Exception: pass seen: set[str] = set() for sheet_name, range_ref in destinations: try: sheet = workbook[sheet_name] resolved_cells = flatten_openpyxl_range_cells(sheet[range_ref]) except Exception: continue for cell in resolved_cells: normalized_value = spreadsheet_text_value(getattr(cell, "value", None)) if not normalized_value: continue key = normalized_value.casefold() if key not in seen: seen.add(key) values.append(normalized_value) if len(values) >= SPREADSHEET_MAX_ENUM_VALUES: return values return values def extract_openpyxl_chart_text(workbook: object, title: object) -> str | None: if title is None: return None if isinstance(title, str): return spreadsheet_text_value(title) or None tx = getattr(title, "tx", None) rich = getattr(tx, "rich", None) or getattr(title, "rich", None) if rich is not None: parts: list[str] = [] for paragraph in getattr(rich, "p", []) or []: for run in getattr(paragraph, "r", []) or []: text_value = spreadsheet_text_value(getattr(run, "t", None)) if text_value: parts.append(text_value) if parts: return " ".join(parts) str_ref = getattr(tx, "strRef", None) or getattr(title, "strRef", None) formula = getattr(str_ref, "f", None) if formula: values = resolve_openpyxl_reference_values(workbook, formula) if values: return values[0] return None def extract_openpyxl_chart_lines(workbook: object, sheet: object) -> list[str]: lines: list[str] = [] for chart in getattr(sheet, "_charts", []) or []: title = extract_openpyxl_chart_text(workbook, getattr(chart, "title", None)) x_axis = extract_openpyxl_chart_text(workbook, getattr(getattr(chart, "x_axis", None), "title", None)) y_axis = extract_openpyxl_chart_text(workbook, getattr(getattr(chart, "y_axis", None), "title", None)) parts: list[str] = [] if title: parts.append(title) if x_axis: parts.append(f"X axis: {x_axis}") if y_axis: parts.append(f"Y axis: {y_axis}") if parts: lines.append(f"- {'; '.join(parts)}") return lines def openpyxl_validation_map(workbook: object, sheet: object, headers: list[str], header_row_index: int | None) -> dict[int, list[str]]: validation_map: dict[int, list[str]] = {} if header_row_index is None: return validation_map data_validations = getattr(getattr(sheet, "data_validations", None), "dataValidation", None) if not data_validations: return validation_map for validation in data_validations: if str(getattr(validation, "type", "")).lower() != "list": continue formula = str(getattr(validation, "formula1", "") or "").strip() if not formula: continue values = resolve_openpyxl_reference_values(workbook, formula) if not values: continue try: ranges = list(getattr(getattr(validation, "cells", None), "ranges", []) or []) except Exception: ranges = [] covered_columns: set[int] = set() for cell_range in ranges: if getattr(cell_range, "max_row", 0) and int(cell_range.max_row) < header_row_index: continue for column_index in range(int(cell_range.min_col), int(cell_range.max_col) + 1): covered_columns.add(column_index) for column_index in sorted(covered_columns): validation_map[column_index] = values[:SPREADSHEET_MAX_ENUM_VALUES] return validation_map def extract_openpyxl_named_ranges(workbook: object) -> list[str]: lines: list[str] = [] for name, defined_name in getattr(workbook, "defined_names", {}).items(): normalized_name = spreadsheet_text_value(name) if not normalized_name: continue lowered = normalized_name.casefold() if lowered in {"_xlnm.print_area", "_xlnm.print_titles"}: continue target = spreadsheet_text_value(getattr(defined_name, "attr_text", None)) if not target: try: destinations = [f"{sheet_name}!{range_ref}" for sheet_name, range_ref in defined_name.destinations] except Exception: destinations = [] target = ", ".join(destinations) if not target: continue lines.append(f"- {normalized_name} -> {target}") if len(lines) >= SPREADSHEET_MAX_NAMED_RANGES: break return lines def scan_openpyxl_sheet_surface( workbook: object, sheet: object, *, read_only: bool, workbook_hyperlinks_seen: set[str], participants: list[str], participant_seen: set[str], ) -> dict[str, object]: output = io.StringIO() writer = csv.writer(output) header_row_index: int | None = None raw_headers: list[object] = [] observed_columns = 0 column_samples: dict[int, list[object]] = defaultdict(list) column_formats: dict[int, list[str]] = defaultdict(list) comment_lines: list[str] = [] comment_overflow = 0 hyperlink_lines: list[str] = [] for row_index, row in enumerate(sheet.iter_rows(), start=1): row_values = [getattr(cell, "value", None) for cell in row] writer.writerow([stringify_spreadsheet_value(value) for value in row_values]) last_nonempty = spreadsheet_row_last_nonempty_index(row_values) if header_row_index is None: if last_nonempty == 0: pass else: header_row_index = row_index raw_headers = row_values[:last_nonempty] observed_columns = max(observed_columns, last_nonempty) else: observed_columns = max(observed_columns, last_nonempty) if last_nonempty > 0: for column_index, cell in enumerate(row[:last_nonempty], start=1): value = getattr(cell, "value", None) if not spreadsheet_value_is_present(value): continue if len(column_samples[column_index]) < SPREADSHEET_TYPE_SAMPLE_LIMIT: column_samples[column_index].append(value) number_format = spreadsheet_text_value(getattr(cell, "number_format", None)) if number_format and len(column_formats[column_index]) < SPREADSHEET_TYPE_SAMPLE_LIMIT: column_formats[column_index].append(number_format) if read_only: continue for cell in row: comment = getattr(cell, "comment", None) if comment is not None: author = spreadsheet_text_value(getattr(comment, "author", None)) or "Unknown" body = normalize_whitespace(str(getattr(comment, "text", "") or "")) if body: append_spreadsheet_participants(participants, participant_seen, [author]) if len(comment_lines) < SPREADSHEET_MAX_COMMENTS_PER_SHEET: comment_lines.append(f"- {author}: {body}") else: comment_overflow += 1 hyperlink = getattr(cell, "hyperlink", None) target = spreadsheet_text_value(getattr(hyperlink, "target", None) or getattr(hyperlink, "location", None)) if not target: continue display = spreadsheet_text_value(getattr(cell, "value", None)) rendered = target if not display or display == target else f"{target} ({display})" hyperlink_key = rendered.casefold() if hyperlink_key in workbook_hyperlinks_seen: continue if len(workbook_hyperlinks_seen) >= SPREADSHEET_MAX_HYPERLINKS_PER_WORKBOOK: continue workbook_hyperlinks_seen.add(hyperlink_key) hyperlink_lines.append(f"- {rendered}") headers = normalize_spreadsheet_headers(raw_headers, observed_columns) validation_map = {} if read_only else openpyxl_validation_map(workbook, sheet, headers, header_row_index) column_descriptors: list[str] = [] visible_column_count = min(len(headers), SPREADSHEET_MAX_COLUMNS_PER_SHEET) for column_index in range(1, visible_column_count + 1): header = headers[column_index - 1] column_type = infer_spreadsheet_column_type( column_samples.get(column_index, []), column_formats.get(column_index, []), force_enum=column_index in validation_map, ) column_descriptors.append(f"{header} ({column_type})") validation_lines: list[str] = [] for column_index in sorted(validation_map): if column_index < 1 or column_index > len(headers): continue values = validation_map[column_index] validation_lines.append(f"- {headers[column_index - 1]} ∈ {{{', '.join(values)}}}") chart_lines = [] if read_only else extract_openpyxl_chart_lines(workbook, sheet) return { "sheet_name": str(sheet.title), "preview_csv": output.getvalue(), "column_descriptors": column_descriptors, "hidden_column_count": max(0, len(headers) - SPREADSHEET_MAX_COLUMNS_PER_SHEET), "comment_lines": comment_lines, "comment_overflow": comment_overflow, "validation_lines": validation_lines, "hyperlink_lines": hyperlink_lines, "chart_lines": chart_lines, } def xls_cell_number_format(workbook: object, cell: object) -> str | None: try: xf = workbook.xf_list[cell.xf_index] format_obj = workbook.format_map.get(xf.format_key) except Exception: return None if format_obj is None: return None return spreadsheet_text_value(getattr(format_obj, "format_str", None)) or None def extract_xls_named_ranges(workbook: object) -> list[str]: lines: list[str] = [] for name_obj in getattr(workbook, "name_obj_list", []) or []: name = spreadsheet_text_value(getattr(name_obj, "name", None)) if not name: continue lowered = name.casefold() if lowered in {"_xlnm.print_area", "_xlnm.print_titles"}: continue target = spreadsheet_text_value(getattr(name_obj, "formula_text", None) or getattr(name_obj, "raw_formula", None)) if not target: continue lines.append(f"- {name} -> {target}") if len(lines) >= SPREADSHEET_MAX_NAMED_RANGES: break return lines def scan_xls_sheet_surface( workbook: object, sheet: object, *, workbook_hyperlinks_seen: set[str], participants: list[str], participant_seen: set[str], ) -> dict[str, object]: output = io.StringIO() writer = csv.writer(output) header_row_index: int | None = None raw_headers: list[object] = [] observed_columns = 0 column_samples: dict[int, list[object]] = defaultdict(list) column_formats: dict[int, list[str]] = defaultdict(list) for row_index in range(sheet.nrows): row = sheet.row(row_index) row_values = [cell.value for cell in row] writer.writerow([stringify_spreadsheet_value(value) for value in row_values]) last_nonempty = spreadsheet_row_last_nonempty_index(row_values) if header_row_index is None: if last_nonempty == 0: continue header_row_index = row_index + 1 raw_headers = row_values[:last_nonempty] observed_columns = max(observed_columns, last_nonempty) continue observed_columns = max(observed_columns, last_nonempty) if last_nonempty == 0: continue for column_index, cell in enumerate(row[:last_nonempty], start=1): if not spreadsheet_value_is_present(cell.value): continue if len(column_samples[column_index]) < SPREADSHEET_TYPE_SAMPLE_LIMIT: column_samples[column_index].append(cell.value) number_format = xls_cell_number_format(workbook, cell) if number_format and len(column_formats[column_index]) < SPREADSHEET_TYPE_SAMPLE_LIMIT: column_formats[column_index].append(number_format) headers = normalize_spreadsheet_headers(raw_headers, observed_columns) column_descriptors: list[str] = [] visible_column_count = min(len(headers), SPREADSHEET_MAX_COLUMNS_PER_SHEET) for column_index in range(1, visible_column_count + 1): header = headers[column_index - 1] column_type = infer_spreadsheet_column_type(column_samples.get(column_index, []), column_formats.get(column_index, [])) column_descriptors.append(f"{header} ({column_type})") comment_lines: list[str] = [] comment_overflow = 0 for note in getattr(sheet, "cell_note_map", {}).values(): author = spreadsheet_text_value(getattr(note, "author", None)) or "Unknown" body = normalize_whitespace(str(getattr(note, "text", "") or "")) if not body: continue append_spreadsheet_participants(participants, participant_seen, [author]) if len(comment_lines) < SPREADSHEET_MAX_COMMENTS_PER_SHEET: comment_lines.append(f"- {author}: {body}") else: comment_overflow += 1 hyperlink_lines: list[str] = [] hyperlink_candidates = list(getattr(sheet, "hyperlink_list", []) or []) hyperlink_candidates.extend(list(getattr(getattr(sheet, "hyperlink_map", {}), "values", lambda: [])())) for hyperlink in hyperlink_candidates: target = spreadsheet_text_value( getattr(hyperlink, "url_or_path", None) or getattr(hyperlink, "target", None) or getattr(hyperlink, "url", None) ) if not target: continue display = spreadsheet_text_value(getattr(hyperlink, "desc", None) or getattr(hyperlink, "description", None)) rendered = target if not display or display == target else f"{target} ({display})" hyperlink_key = rendered.casefold() if hyperlink_key in workbook_hyperlinks_seen: continue if len(workbook_hyperlinks_seen) >= SPREADSHEET_MAX_HYPERLINKS_PER_WORKBOOK: continue workbook_hyperlinks_seen.add(hyperlink_key) hyperlink_lines.append(f"- {rendered}") return { "sheet_name": str(sheet.name), "preview_csv": output.getvalue(), "column_descriptors": column_descriptors, "hidden_column_count": max(0, len(headers) - SPREADSHEET_MAX_COLUMNS_PER_SHEET), "comment_lines": comment_lines, "comment_overflow": comment_overflow, "validation_lines": [], "hyperlink_lines": hyperlink_lines, "chart_lines": [], } def extract_xlsx_file(path: Path) -> dict[str, object]: openpyxl_module = dependency_guard("openpyxl", "openpyxl", "xlsx") read_only = path.stat().st_size > SPREADSHEET_XLSX_READ_ONLY_FALLBACK_BYTES workbook = openpyxl_module.load_workbook(path, read_only=read_only, data_only=True) # type: ignore[union-attr] preview_artifacts: list[dict[str, object]] = [build_spreadsheet_html_preview_artifact(path, ordinal=0)] sections: list[str] = [] workbook_hyperlinks_seen: set[str] = set() participants: list[str] = [] participant_seen: set[str] = set() try: props = getattr(workbook, "properties", None) author = spreadsheet_text_value(getattr(props, "creator", None)) or None title = spreadsheet_text_value(getattr(props, "title", None)) or path.stem subject = spreadsheet_text_value(getattr(props, "subject", None)) or None last_modified_by = spreadsheet_text_value(getattr(props, "lastModifiedBy", None)) or None append_spreadsheet_participants(participants, participant_seen, [author, last_modified_by]) sheet_names = [str(sheet.title) for sheet in workbook.worksheets] sheet_count = len(sheet_names) sections.append( build_spreadsheet_header_section( title=title, sheet_names=sheet_names, author=author, subject=subject, participants=None, parse_note=( "Parse note: large-workbook read-only fallback; workbook properties and named ranges are preserved, " "but comments, hyperlinks, validations, charts, and other sheet-level details may be incomplete." if read_only else None ), ) ) for ordinal, sheet in enumerate(workbook.worksheets, start=1): surface = scan_openpyxl_sheet_surface( workbook, sheet, read_only=read_only, workbook_hyperlinks_seen=workbook_hyperlinks_seen, participants=participants, participant_seen=participant_seen, ) preview_artifacts.append( { "file_name": f"{path.name}.{slugify(str(surface['sheet_name']))}.csv", "preview_type": "csv", "label": surface["sheet_name"], "ordinal": ordinal, "content": surface["preview_csv"], } ) sections.append( build_spreadsheet_sheet_section( str(surface["sheet_name"]), list(surface["column_descriptors"]), hidden_column_count=int(surface["hidden_column_count"]), comment_lines=list(surface["comment_lines"]), comment_overflow=int(surface["comment_overflow"]), validation_lines=list(surface["validation_lines"]), hyperlink_lines=list(surface["hyperlink_lines"]), chart_lines=list(surface["chart_lines"]), ) ) named_ranges = extract_openpyxl_named_ranges(workbook) if named_ranges: named_range_section_lines = ["Named ranges:"] named_range_section_lines.extend(named_ranges) if len(named_ranges) >= SPREADSHEET_MAX_NAMED_RANGES: named_range_section_lines.append("Named ranges truncated: more named ranges not listed.") sections.append(normalize_whitespace("\n".join(named_range_section_lines))) finally: workbook.close() participants_text = ", ".join(participants) or None if participants_text: sections[0] = build_spreadsheet_header_section( title=title, sheet_names=sheet_names, author=author, subject=subject, participants=participants_text, parse_note=( "Parse note: large-workbook read-only fallback; workbook properties and named ranges are preserved, " "but comments, hyperlinks, validations, charts, and other sheet-level details may be incomplete." if read_only else None ), ) text_content, chunks = build_structural_summary_from_sections(sections) return { "page_count": sheet_count, "author": author, "content_type": "Spreadsheet / Table", "date_created": normalize_datetime(getattr(props, "created", None)), "date_modified": normalize_datetime(getattr(props, "modified", None)), "participants": participants_text, "title": title, "subject": subject, "recipients": None, "text_content": text_content, "text_status": "empty" if not text_content else "ok", "chunks": chunks, "preview_artifacts": preview_artifacts, } def extract_xls_file(path: Path) -> dict[str, object]: xlrd_module = dependency_guard("xlrd", "xlrd", "xls") workbook = xlrd_module.open_workbook(path, formatting_info=True) # type: ignore[union-attr] preview_artifacts: list[dict[str, object]] = [build_spreadsheet_html_preview_artifact(path, ordinal=0)] sections: list[str] = [] workbook_hyperlinks_seen: set[str] = set() participants: list[str] = [] participant_seen: set[str] = set() author = spreadsheet_text_value(getattr(workbook, "user_name", None)) or None append_spreadsheet_participants(participants, participant_seen, [author]) sheet_names = [str(sheet.name) for sheet in workbook.sheets()] sheet_count = len(sheet_names) sections.append( build_spreadsheet_header_section( title=path.stem, sheet_names=sheet_names, author=author, subject=None, participants=None, ) ) for ordinal, sheet in enumerate(workbook.sheets(), start=1): surface = scan_xls_sheet_surface( workbook, sheet, workbook_hyperlinks_seen=workbook_hyperlinks_seen, participants=participants, participant_seen=participant_seen, ) preview_artifacts.append( { "file_name": f"{path.name}.{slugify(str(surface['sheet_name']))}.csv", "preview_type": "csv", "label": surface["sheet_name"], "ordinal": ordinal, "content": surface["preview_csv"], } ) sections.append( build_spreadsheet_sheet_section( str(surface["sheet_name"]), list(surface["column_descriptors"]), hidden_column_count=int(surface["hidden_column_count"]), comment_lines=list(surface["comment_lines"]), comment_overflow=int(surface["comment_overflow"]), validation_lines=list(surface["validation_lines"]), hyperlink_lines=list(surface["hyperlink_lines"]), chart_lines=list(surface["chart_lines"]), ) ) named_ranges = extract_xls_named_ranges(workbook) if named_ranges: named_range_section_lines = ["Named ranges:"] named_range_section_lines.extend(named_ranges) if len(named_ranges) >= SPREADSHEET_MAX_NAMED_RANGES: named_range_section_lines.append("Named ranges truncated: more named ranges not listed.") sections.append(normalize_whitespace("\n".join(named_range_section_lines))) participants_text = ", ".join(participants) or None if participants_text: sections[0] = build_spreadsheet_header_section( title=path.stem, sheet_names=sheet_names, author=author, subject=None, participants=participants_text, ) text_content, chunks = build_structural_summary_from_sections(sections) return { "page_count": sheet_count, "author": author, "content_type": "Spreadsheet / Table", "date_created": None, "date_modified": None, "participants": participants_text, "title": path.stem, "subject": None, "recipients": None, "text_content": text_content, "text_status": "empty" if not text_content else "ok", "chunks": chunks, "preview_artifacts": preview_artifacts, } def parse_delimited_rows_for_summary(path: Path) -> tuple[list[list[str]], str]: decoded, text_status, _ = decode_bytes(path.read_bytes()) normalized_text = decoded.lstrip("\ufeff") rows_source = normalized_text delimiter: str | None = "\t" if normalize_extension(path) == "tsv" else None first_line, _, remainder = normalized_text.partition("\n") directive_match = re.fullmatch(r"\s*sep=(.)\s*", first_line) if directive_match is not None: delimiter = directive_match.group(1) rows_source = remainder if delimiter is None: sample = rows_source[:4096] try: dialect = csv.Sniffer().sniff(sample, delimiters=",|\t;") reader = csv.reader(io.StringIO(rows_source), dialect) except csv.Error: reader = csv.reader(io.StringIO(rows_source), csv.excel) else: reader = csv.reader(io.StringIO(rows_source), delimiter=delimiter) rows = [[field.strip().strip("\ufeff") for field in row] for row in reader if any(field.strip() for field in row)] return rows, text_status def csv_first_row_looks_like_header(row: list[str]) -> bool: nonempty = [field.strip() for field in row if field.strip()] if not nonempty: return False if len({value.casefold() for value in nonempty}) != len(nonempty): return False if any(len(value) > 120 for value in nonempty): return False numeric_like = sum(1 for value in nonempty if spreadsheet_string_shape(value) in {"integer", "number", "currency", "percent", "date", "datetime"}) return numeric_like < len(nonempty) def extract_csv_file(path: Path) -> dict[str, object]: rows, text_status = parse_delimited_rows_for_summary(path) title = path.stem preview_artifacts = [build_spreadsheet_html_preview_artifact(path, ordinal=0)] if not rows: sections = [ build_spreadsheet_header_section( title=title, sheet_names=[title], author=None, subject=None, participants=None, ), build_spreadsheet_sheet_section(title, []), ] text_content, chunks = build_structural_summary_from_sections(sections) return { "page_count": 1, "author": None, "content_type": "Spreadsheet / Table", "date_created": None, "date_modified": None, "participants": None, "title": title, "subject": None, "recipients": None, "text_content": text_content, "text_status": "empty" if not text_content else text_status, "chunks": chunks, "preview_artifacts": preview_artifacts, } has_headers = csv_first_row_looks_like_header(rows[0]) header_row = rows[0] if has_headers else [] data_rows = rows[1:] if has_headers else rows column_count = max((len(row) for row in rows), default=0) headers = normalize_spreadsheet_headers(header_row, column_count) column_samples: dict[int, list[object]] = defaultdict(list) for row in data_rows: for column_index in range(1, min(len(row), column_count) + 1): value = row[column_index - 1] if not spreadsheet_value_is_present(value): continue if len(column_samples[column_index]) < SPREADSHEET_TYPE_SAMPLE_LIMIT: column_samples[column_index].append(value) column_descriptors: list[str] = [] visible_column_count = min(len(headers), SPREADSHEET_MAX_COLUMNS_PER_SHEET) for column_index in range(1, visible_column_count + 1): column_descriptors.append( f"{headers[column_index - 1]} ({infer_spreadsheet_column_type(column_samples.get(column_index, []), [])})" ) sections = [ build_spreadsheet_header_section( title=title, sheet_names=[title], author=None, subject=None, participants=None, ), build_spreadsheet_sheet_section( title, column_descriptors, hidden_column_count=max(0, len(headers) - SPREADSHEET_MAX_COLUMNS_PER_SHEET), ), ] text_content, chunks = build_structural_summary_from_sections(sections) return { "page_count": 1, "author": None, "content_type": "Spreadsheet / Table", "date_created": None, "date_modified": None, "participants": None, "title": title, "subject": None, "recipients": None, "text_content": text_content, "text_status": "empty" if not text_content else text_status, "chunks": chunks, "preview_artifacts": preview_artifacts, } def normalize_attachment_filename( file_name: str | None, ordinal: int, *, payload: bytes | None = None, content_type: object = None, preferred_extension: object = None, ) -> str: normalized = normalize_whitespace(file_name or "") detected_file_type = infer_attachment_file_type( file_name=normalized or None, payload=payload, content_type=content_type, preferred_extension=preferred_extension, ) if normalized: current_file_type = normalize_file_type_name(Path(normalized).suffix.lower().lstrip(".")) if ( current_file_type == "bin" and re.fullmatch(rf"attachment-{ordinal:03d}\.bin", normalized, re.IGNORECASE) and detected_file_type and detected_file_type != "bin" ): return f"attachment-{ordinal:03d}.{detected_file_type}" if current_file_type: return normalized if detected_file_type: return f"{normalized}.{detected_file_type}" return normalized return f"attachment-{ordinal:03d}.{detected_file_type or 'bin'}" def coerce_email_part_payload(part: object) -> bytes | None: try: payload = part.get_payload(decode=True) except Exception: payload = None if isinstance(payload, bytes): return payload try: content = part.get_content() except Exception: content = None if isinstance(content, bytes): return content if isinstance(content, str): charset = None try: charset = part.get_content_charset() except Exception: charset = None return content.encode(charset or "utf-8", errors="replace") if hasattr(content, "as_bytes"): try: return content.as_bytes(policy=policy.default) except Exception: return None return None def extract_eml_attachments(message: object) -> list[dict[str, object]]: attachments: list[dict[str, object]] = [] ordinal = 1 for part in message.walk(): if part.is_multipart(): continue disposition = (part.get_content_disposition() or "").lower() file_name = part.get_filename() content_type = normalize_mime_type(part.get_content_type()) content_id = normalize_content_id(part.get("Content-ID")) is_inline = disposition == "inline" or (content_id is not None and disposition != "attachment") if disposition != "attachment" and not file_name and not content_id: continue payload = coerce_email_part_payload(part) if payload is None: continue attachments.append( { "file_name": normalize_attachment_filename( file_name, ordinal, payload=payload, content_type=content_type, ), "ordinal": ordinal, "payload": payload, "file_hash": sha256_bytes(payload), "content_type": content_type, "content_id": content_id, "is_inline": is_inline, } ) ordinal += 1 return attachments def extract_msg_attachment_payload(attachment: object) -> bytes | None: data = getattr(attachment, "data", None) if isinstance(data, bytes): return data if isinstance(data, str): return data.encode("utf-8", errors="replace") if hasattr(data, "as_bytes"): try: return data.as_bytes(policy=policy.default) except Exception: return None try: with tempfile.TemporaryDirectory(prefix="retriever-msg-attachment-") as tempdir: saved = attachment.save(customPath=tempdir, extractEmbedded=True) if isinstance(saved, tuple) and len(saved) >= 2 and saved[1]: saved_path = Path(str(saved[1])) if saved_path.exists(): return saved_path.read_bytes() except Exception: return None return None def extract_msg_attachment_content_id(attachment: object) -> str | None: for attr_name in ("cid", "contentId", "content_id"): raw = getattr(attachment, attr_name, None) normalized = normalize_content_id(raw) if normalized: return normalized return None def extract_msg_attachment_content_type(attachment: object) -> str | None: for attr_name in ("mimeTag", "mime_tag", "mimetype", "mime_type", "contentType", "content_type"): normalized = normalize_mime_type(getattr(attachment, attr_name, None)) if normalized: return normalized return None def extract_msg_attachments(message: object) -> list[dict[str, object]]: attachments: list[dict[str, object]] = [] ordinal = 1 for attachment in getattr(message, "attachments", []): if getattr(attachment, "hidden", False): continue try: file_name = attachment.getFilename() except Exception: file_name = None content_type = extract_msg_attachment_content_type(attachment) payload = extract_msg_attachment_payload(attachment) if payload is None: continue attachments.append( { "file_name": normalize_attachment_filename( file_name, ordinal, payload=payload, content_type=content_type, ), "ordinal": ordinal, "payload": payload, "file_hash": sha256_bytes(payload), "content_type": content_type, "content_id": extract_msg_attachment_content_id(attachment), } ) ordinal += 1 return attachments def build_email_extracted_payload( *, subject: str | None, author: str | None, recipients: str | None, date_created: str | None, text_body: str | None, html_body: str | None, attachments: list[dict[str, object]] | None, preview_file_name: str, email_threading: dict[str, object] | None = None, ) -> dict[str, object]: normalized_html = None if html_body is None else str(html_body) normalized_text = normalize_whitespace(str(text_body or "")) if not normalized_text and normalized_html: normalized_text = strip_html_tags(normalized_html) normalized_text = normalize_whitespace(normalized_text) resolved_subject = normalize_generated_document_title(subject) participants = extract_email_chain_participants( normalized_text, [author, recipients or None], ) normalized_attachments = list(attachments or []) preview_html_body = inline_cid_references_in_html(normalized_html, normalized_attachments) document_attachments = filter_html_preview_embedded_image_attachments( normalized_html, normalized_attachments, ) calendar_invites, document_attachments = partition_calendar_invite_attachments(document_attachments) if calendar_invites: merged_participants: list[str] = [] seen_participants: set[str] = set() append_unique_participants(merged_participants, seen_participants, [participants]) for invite in calendar_invites: append_unique_participants( merged_participants, seen_participants, [invite.get("organizer"), invite.get("attendees")], ) participants = ", ".join(merged_participants) or None calendar_invite_search_text = build_calendar_invite_search_text(calendar_invites) indexed_text_content = "\n\n".join( part for part in (normalized_text, calendar_invite_search_text) if part ) preview = build_email_message_preview_html( { "id": 0, "author": author, "recipients": recipients or None, "date_created": date_created, "subject": resolved_subject, "title": resolved_subject, "text_content": indexed_text_content, "standalone_preview_body_html": preview_html_body, }, body_html=preview_html_body, ) return { "page_count": 1, "author": author, "content_type": "Email", "date_created": date_created, "date_modified": None, "participants": participants, "title": resolved_subject, "subject": resolved_subject, "recipients": recipients or None, "text_content": indexed_text_content, "text_status": "empty" if not indexed_text_content else "ok", "attachments": document_attachments, "email_threading": dict(email_threading or {}), "preview_artifacts": [ { "file_name": preview_file_name, "preview_type": "html", "label": "message", "ordinal": 0, "content": preview, } ], } def parse_email_headers_only(header_text: object) -> object | None: normalized = str(header_text or "") if not normalize_whitespace(normalized): return None payload = normalized.replace("\r\n", "\n").replace("\r", "\n") if "\n\n" not in payload: payload = payload.rstrip("\n") + "\n\n" try: return BytesParser(policy=policy.default).parsebytes(payload.encode("utf-8", errors="replace"), headersonly=True) except Exception: return None def email_header_value(header_mapping: object, key: str) -> str | None: target = normalize_whitespace(key).lower() if isinstance(header_mapping, dict): for raw_key, raw_value in header_mapping.items(): if normalize_whitespace(str(raw_key or "")).lower() != target: continue normalized = normalize_whitespace(str(raw_value or "")) return normalized or None return None def extract_email_recipients_from_headers(header_text: object) -> str | None: parsed_headers = parse_email_headers_only(header_text) if parsed_headers is None: return None recipients: list[str] = [] seen: set[str] = set() append_unique_participants( recipients, seen, [ normalize_whitespace(str(parsed_headers.get("To") or "")) or None, normalize_whitespace(str(parsed_headers.get("Cc") or "")) or None, normalize_whitespace(str(parsed_headers.get("Bcc") or "")) or None, ], ) return ", ".join(recipients) or None def build_email_threading_payload( *, subject: object, message_id: object = None, in_reply_to: object = None, references: object = None, conversation_index: object = None, conversation_topic: object = None, ) -> dict[str, object]: normalized_conversation_topic = normalize_email_thread_subject(conversation_topic) normalized_subject = normalize_email_thread_subject(subject or conversation_topic) return { "message_id": normalize_email_message_id(message_id), "in_reply_to": normalize_email_message_id(in_reply_to), "references": extract_email_message_ids(references), "conversation_index": normalize_whitespace(str(conversation_index or "")) or None, "conversation_topic": normalized_conversation_topic, "normalized_subject": normalized_subject, } def extract_parsed_email_threading(message: object) -> dict[str, object]: return build_email_threading_payload( subject=message.get("Subject"), message_id=message.get("Message-ID") or message.get("Message-Id"), in_reply_to=message.get("In-Reply-To"), references=message.get("References"), conversation_index=message.get("Conversation-Index"), conversation_topic=message.get("Conversation-Topic"), ) def extract_msg_threading(message: object, subject: object) -> dict[str, object]: header_mapping = getattr(message, "headerDict", None) parsed_headers = parse_email_headers_only( getattr(message, "headerText", None) or getattr(message, "header", None) ) def _value(header_name: str) -> str | None: mapped = email_header_value(header_mapping, header_name) if mapped: return mapped if parsed_headers is not None: normalized = normalize_whitespace(str(parsed_headers.get(header_name) or "")) if normalized: return normalized return None return build_email_threading_payload( subject=subject, message_id=_value("Message-ID"), in_reply_to=_value("In-Reply-To"), references=_value("References"), conversation_index=_value("Conversation-Index"), conversation_topic=_value("Conversation-Topic"), ) def extract_transport_header_threading( transport_headers: object, *, subject: object, conversation_topic: object = None, ) -> dict[str, object]: parsed_headers = parse_email_headers_only(transport_headers) return build_email_threading_payload( subject=subject, message_id=parsed_headers.get("Message-ID") if parsed_headers is not None else None, in_reply_to=parsed_headers.get("In-Reply-To") if parsed_headers is not None else None, references=parsed_headers.get("References") if parsed_headers is not None else None, conversation_index=parsed_headers.get("Conversation-Index") if parsed_headers is not None else None, conversation_topic=conversation_topic or (parsed_headers.get("Conversation-Topic") if parsed_headers is not None else None), ) def build_chat_preview_artifacts( *, title: str | None, author: str | None, participants: str | None, date_created: str | None, date_modified: str | None, text_body: str | None, preview_file_name: str, chat_metadata: dict[str, object] | None = None, chat_entries: list[dict[str, object]] | None = None, label: str = "conversation", ) -> list[dict[str, object]]: normalized_text = normalize_whitespace(str(text_body or "")) metadata = chat_metadata or extract_chat_transcript_metadata(normalized_text) or {} headers = { "Author": author or "", "Participants": participants or "", "Started": format_chat_preview_timestamp(date_created) or date_created or "", "Updated": format_chat_preview_timestamp(date_modified) or date_modified or "", "Title": title or "", "Messages": str(metadata["message_count"]) if metadata.get("message_count") is not None else "", } preview = build_chat_preview_html( headers, normalized_text, document_title=title or "Retriever Chat Preview", entries=chat_entries, ) return [ { "file_name": preview_file_name, "preview_type": "html", "label": label, "ordinal": 0, "content": preview, } ] def build_chat_extracted_payload( *, title: str | None, author: str | None, date_created: str | None, text_body: str | None, html_body: str | None, attachments: list[dict[str, object]] | None, preview_file_name: str, chat_metadata: dict[str, object] | None = None, chat_entries: list[dict[str, object]] | None = None, chat_threading: dict[str, object] | None = None, ) -> dict[str, object]: normalized_html = None if html_body is None else str(html_body) normalized_text = normalize_whitespace(str(text_body or "")) if not normalized_text and normalized_html: normalized_text = strip_html_tags(normalized_html) normalized_text = normalize_whitespace(normalized_text) metadata = chat_metadata or extract_chat_transcript_metadata(normalized_text) or {} participants = ( str(metadata["participants"]) if metadata.get("participants") else extract_chat_participants(normalized_text) ) resolved_date_created = str(metadata["date_created"]) if metadata.get("date_created") else date_created resolved_date_modified = str(metadata["date_modified"]) if metadata.get("date_modified") else None metadata_title = normalize_generated_document_title(str(metadata["title"])) if metadata.get("title") else None resolved_title = normalize_generated_document_title(title) or metadata_title return { "page_count": 1, "author": None, "content_type": "Chat", "date_created": resolved_date_created, "date_modified": resolved_date_modified, "participants": participants, "title": resolved_title, "subject": None, "recipients": None, "text_content": normalized_text, "text_status": "empty" if not normalized_text else "ok", "chat_threading": dict(chat_threading or {}), "attachments": list(attachments or []), "preview_artifacts": build_chat_preview_artifacts( title=resolved_title, author=None, participants=participants, date_created=resolved_date_created, date_modified=resolved_date_modified, text_body=normalized_text, preview_file_name=preview_file_name, chat_metadata=metadata, chat_entries=chat_entries, ), } def build_calendar_extracted_payload( *, subject: str | None, author: str | None, recipients: str | None, date_created: str | None, text_body: str | None, html_body: str | None, attachments: list[dict[str, object]] | None, preview_file_name: str, ) -> dict[str, object]: normalized_html = None if html_body is None else str(html_body) normalized_text = normalize_whitespace(str(text_body or "")) if not normalized_text and normalized_html: normalized_text = strip_html_tags(normalized_html) normalized_text = normalize_whitespace(normalized_text) resolved_subject = normalize_generated_document_title(subject) normalized_attachments = list(attachments or []) preview_html_body = inline_cid_references_in_html(normalized_html, normalized_attachments) preview_title = resolved_subject or "Retriever Calendar Preview" preview = build_html_preview( { "Organizer": author or "", "Date": format_chat_preview_timestamp(date_created) or date_created or "", "Title": resolved_subject or "", "Attendees": recipients or "", }, body_html=preview_html_body, body_text=normalized_text, document_title=preview_title, ) return { "page_count": 1, "author": author, "content_type": "Calendar", "date_created": date_created, "date_modified": None, "participants": author or None, "title": resolved_subject, "subject": resolved_subject, "recipients": recipients or None, "text_content": normalized_text, "text_status": "empty" if not normalized_text else "ok", "attachments": filter_html_preview_embedded_image_attachments( normalized_html, normalized_attachments, ), "preview_artifacts": [ { "file_name": preview_file_name, "preview_type": "html", "label": "calendar", "ordinal": 0, "content": preview, } ], } def normalize_pst_folder_marker(folder_path: object) -> str: normalized = normalize_whitespace(str(folder_path or "")).strip().strip("/") if not normalized: return "/" return f"/{normalized.lower()}/" def pst_folder_path_contains(folder_path: object, marker: str) -> bool: return marker in normalize_pst_folder_marker(folder_path) def infer_pst_chat_title(subject: str | None, text_body: str | None) -> str | None: normalized_subject = normalize_generated_document_title(subject) if normalized_subject: return normalized_subject normalized_text = normalize_whitespace(str(text_body or "")) if not normalized_text: return None if len(normalized_text) <= 120: return normalized_text return normalized_text[:117].rstrip() + "..." def synthesize_pst_chat_entries( *, author: str | None, date_created: str | None, text_body: str | None, chat_metadata: dict[str, object] | None, ) -> list[dict[str, object]] | None: if chat_metadata: try: if int(chat_metadata.get("message_count") or 0) > 1: return None except Exception: pass speaker = normalize_whitespace(str(author or "")) body = normalize_whitespace(str(text_body or "")) if not speaker or not body: return None return [ { "speaker": speaker, "body": body, "timestamp": date_created, "timestamp_label": format_chat_preview_timestamp(date_created), "avatar_label": chat_avatar_initials(speaker), } ] def synthesize_pst_chat_metadata( *, subject: str | None, author: str | None, date_created: str | None, text_body: str | None, chat_metadata: dict[str, object] | None, chat_entries: list[dict[str, object]] | None, preferred_title: str | None = None, preferred_participants: str | None = None, ) -> dict[str, object] | None: metadata = dict(chat_metadata or {}) resolved_title = infer_pst_chat_title(subject, text_body) if preferred_title: metadata["title"] = preferred_title elif resolved_title and not metadata.get("title"): metadata["title"] = resolved_title normalized_author = normalize_whitespace(str(author or "")) if preferred_participants: metadata["participants"] = preferred_participants elif normalized_author and not metadata.get("participants"): metadata["participants"] = normalized_author if date_created and not metadata.get("date_created"): metadata["date_created"] = date_created if date_created and not metadata.get("date_modified") and chat_entries is not None: metadata["date_modified"] = date_created if chat_entries is not None and not metadata.get("message_count"): metadata["message_count"] = len(chat_entries) return metadata or None def classify_pst_message_kind(message_dict: dict[str, object], chat_metadata: dict[str, object] | None) -> str: folder_path = message_dict.get("folder_path") if pst_folder_path_contains(folder_path, "/substratefiles/spools/"): return "skip" if pst_folder_path_contains(folder_path, "/skypespacesdata/teamsmeetings/"): return "skip" if pst_folder_path_contains(folder_path, "/teamsmessagesdata/"): return "chat" if pst_folder_path_contains(folder_path, "/conversation history/"): return "chat" message_class = normalize_whitespace( str(message_dict.get("message_class") or message_dict.get("item_class") or "") ) lowered_class = message_class.lower() if message_class else "" if lowered_class.startswith("ipm.appointment") or "appointment" in lowered_class: return "calendar" if any(token in lowered_class for token in ("conversation", "chat", "im")): return "chat" if pst_folder_path_contains(folder_path, "/calendar/"): return "calendar" if not chat_metadata: return "email" recipients = normalize_whitespace(str(message_dict.get("recipients") or "")) return "chat" if not recipients else "email" def extract_parsed_email_message( message: object, *, include_attachments: bool = True, preview_file_name: str = "message.html", ) -> dict[str, object]: subject = message.get("Subject") author = message.get("From") recipients = ", ".join(part for part in [message.get("To"), message.get("Cc"), message.get("Bcc")] if part) date_created = normalize_datetime(message.get("Date")) plain_body = message.get_body(preferencelist=("plain",)) html_body = message.get_body(preferencelist=("html",)) text_body = None html_content = None if plain_body is not None: content = plain_body.get_content() text_body = normalize_whitespace(str(content)) if html_body is not None: content = html_body.get_content() html_content = str(content) if text_body is None: text_body = strip_html_tags(html_content) if text_body is None: content = message.get_content() if isinstance(content, bytes): text_body, _, _ = decode_bytes(content) else: text_body = normalize_whitespace(str(content)) email_threading = extract_parsed_email_threading(message) return build_email_extracted_payload( subject=subject or None, author=author or None, recipients=recipients or None, date_created=date_created, text_body=text_body, html_body=html_content, attachments=extract_eml_attachments(message) if include_attachments else [], preview_file_name=preview_file_name, email_threading=email_threading, ) def parse_email_message( data: bytes, *, include_attachments: bool = True, preview_file_name: str = "message.html", ) -> dict[str, object]: message = BytesParser(policy=policy.default).parsebytes(data) return extract_parsed_email_message( message, include_attachments=include_attachments, preview_file_name=preview_file_name, ) def extract_eml_file(path: Path, include_attachments: bool = True) -> dict[str, object]: return parse_email_message( path.read_bytes(), include_attachments=include_attachments, preview_file_name=f"{path.name}.html", ) def extract_msg_file(path: Path, include_attachments: bool = True) -> dict[str, object]: extract_msg_module = dependency_guard("extract_msg", "extract-msg", "msg") message = extract_msg_module.Message(str(path)) # type: ignore[union-attr] try: subject = message.subject or None author = message.sender or None recipients = ", ".join(part for part in [message.to, message.cc, message.bcc] if part) html_body = None if getattr(message, "htmlBody", None): body_value = message.htmlBody html_body = body_value.decode("utf-8", errors="replace") if isinstance(body_value, bytes) else str(body_value) text_body = normalize_whitespace(str(message.body or "")) or (strip_html_tags(html_body) if html_body else "") email_threading = extract_msg_threading(message, subject) return build_email_extracted_payload( subject=subject, author=author, recipients=recipients or None, date_created=normalize_datetime(message.date), text_body=text_body, html_body=html_body, attachments=extract_msg_attachments(message) if include_attachments else [], preview_file_name=f"{path.name}.html", email_threading=email_threading, ) finally: try: message.close() except Exception: pass def iter_mbox_messages(path: Path): archive = mailbox.mbox(str(path), factory=mailbox.mboxMessage, create=False) duplicate_counts: dict[str, int] = defaultdict(int) try: for _, raw_message in archive.iteritems(): payload_bytes = raw_message.as_bytes(policy=policy.default, unixfrom=False) payload_hash = sha256_bytes(payload_bytes) parsed_message = BytesParser(policy=policy.default).parsebytes(payload_bytes) explicit_source_item_id = normalize_whitespace( str(parsed_message.get("Message-ID") or parsed_message.get("Message-Id") or "") ) or None base_source_item_id = explicit_source_item_id or f"mbox-hash:{payload_hash}" duplicate_counts[base_source_item_id] += 1 occurrence = duplicate_counts[base_source_item_id] stable_source_item_id = base_source_item_id if occurrence == 1 else f"{base_source_item_id}#{occurrence}" yield { "source_item_id": stable_source_item_id, "payload_hash": payload_hash, "parsed_message": parsed_message, } finally: try: archive.close() except Exception: pass def container_message_payload_hash(message_payload: dict[str, object]) -> str: attachment_manifest: list[dict[str, object]] = [] for ordinal, attachment in enumerate(list(message_payload.get("attachments") or []), start=1): raw_name = attachment.get("file_name") if isinstance(attachment, dict) else None payload = attachment.get("payload") if isinstance(attachment, dict) else None file_hash = attachment.get("file_hash") if isinstance(attachment, dict) else None if file_hash is None and isinstance(payload, bytes): file_hash = sha256_bytes(payload) attachment_manifest.append( { "file_name": normalize_attachment_filename(raw_name if isinstance(raw_name, str) else None, ordinal), "ordinal": int(attachment.get("ordinal", ordinal)) if isinstance(attachment, dict) else ordinal, "file_hash": file_hash, "payload_size": len(payload) if isinstance(payload, bytes) else None, } ) return sha256_json_value( { "subject": message_payload.get("subject"), "author": message_payload.get("author"), "recipients": message_payload.get("recipients"), "date_created": message_payload.get("date_created"), "text_body": normalize_whitespace(str(message_payload.get("text_body") or "")), "html_body": normalize_whitespace(strip_html_tags(str(message_payload.get("html_body") or ""))), "attachments": attachment_manifest, } ) def normalize_mbox_message(source_rel_path: str, message_dict: dict[str, object]) -> dict[str, object]: parsed_message = message_dict.get("parsed_message") if parsed_message is None: raise RetrieverError(f"MBOX message is missing a parsed message payload in {source_rel_path}") source_item_id = normalize_source_item_id(message_dict.get("source_item_id") or parsed_message.get("Message-ID")) extracted = extract_parsed_email_message( parsed_message, include_attachments=True, preview_file_name=mbox_preview_file_name(source_item_id), ) payload_hash = normalize_whitespace(str(message_dict.get("payload_hash") or "")) or None return { "rel_path": mbox_message_rel_path(source_rel_path, source_item_id), "file_name": mbox_message_file_name(source_item_id), "file_hash": payload_hash or container_message_payload_hash( { "subject": extracted.get("subject"), "author": extracted.get("author"), "recipients": extracted.get("recipients"), "date_created": extracted.get("date_created"), "text_body": extracted.get("text_content"), "attachments": extracted.get("attachments"), } ), "source_rel_path": source_rel_path, "source_item_id": source_item_id, "source_folder_path": None, "extracted": extracted, } def coerce_pst_attachment_payload(attachment: object) -> bytes | None: if attachment is None: return None if isinstance(attachment, bytes): return attachment if isinstance(attachment, str): return attachment.encode("utf-8", errors="replace") if isinstance(attachment, dict): for key in ("payload", "data", "content"): if key in attachment: return coerce_pst_attachment_payload(attachment.get(key)) for attr_name in ("data", "payload", "content"): value = getattr(attachment, attr_name, None) if isinstance(value, bytes): return value if isinstance(value, str): return value.encode("utf-8", errors="replace") for method_name in ("read_buffer", "read_data", "get_data"): method = getattr(attachment, method_name, None) if not callable(method): continue try: size = getattr(attachment, "size", None) if method_name == "read_buffer" and isinstance(size, int) and size > 0: payload = method(size) else: payload = method() except Exception: continue if isinstance(payload, bytes): return payload if isinstance(payload, str): return payload.encode("utf-8", errors="replace") return None PST_PROP_DISPLAY_NAME = 0x3001 PST_PROP_ATTACH_EXTENSION = 0x3703 PST_PROP_ATTACH_FILENAME = 0x3704 PST_PROP_ATTACH_LONG_FILENAME = 0x3707 PST_PROP_ATTACH_MIME_TAG = 0x370E PST_PROP_ATTACH_CONTENT_ID = 0x3712 PST_SENDER_DISPLAY_NAME_ENTRY_TYPES = {0x0042, 0x0C1A} PST_SENDER_EMAIL_ADDRESS_ENTRY_TYPES = {0x0065, 0x0C1F} PST_MAPI_AUTHOR_DISPLAY_NAME_ENTRY_TYPES = {0x0042, 0x0C1A} PST_MAPI_AUTHOR_SMTP_ADDRESS_ENTRY_TYPES = {0x5D01, 0x5D02} PST_MAPI_AUTHOR_LEGACY_DN_ENTRY_TYPES = {0x0065, 0x0C1F} PST_MAPI_RECIPIENT_DISPLAY_NAME_ENTRY_TYPES = {0x3FF8, 0x3FFA, 0x4038, 0x4039} PST_MAPI_RECIPIENT_SMTP_ADDRESS_ENTRY_TYPES = {0x5D0A, 0x5D0B} PST_MAPI_RECIPIENT_LEGACY_DN_ENTRY_TYPES = {0x4023, 0x4025} PST_TEAMS_ORGID_PREFIX = "8:orgid:" PST_DASHED_GUID_PATTERN = re.compile( r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", re.IGNORECASE, ) PST_DEBUG_SCOPE_NAME_PATTERN = re.compile( r"(team|teams|skype|conversation|thread|chat|channel|space|group|participant|member|roster)", re.IGNORECASE, ) PST_DEBUG_GUID_VALUE_PATTERN = re.compile( r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", re.IGNORECASE, ) PST_DEBUG_MAX_TEXT_VALUE_CHARS = 512 PST_DEBUG_MAX_HEX_PREVIEW_BYTES = 32 def normalize_pst_identifier(value: object) -> str | None: if value is None: return None if isinstance(value, memoryview): value = bytes(value) if isinstance(value, bytes): return value.hex() if isinstance(value, (int, float)): return str(int(value)) normalized = normalize_whitespace(str(value)) return normalized or None def iter_pst_collection( owner: object, *, list_attrs: tuple[str, ...], count_getter_pairs: tuple[tuple[str, str], ...], ): for attr_name in list_attrs: value = getattr(owner, attr_name, None) if value is None: continue try: for item in value: yield item return except TypeError: pass for count_name, getter_name in count_getter_pairs: try: count = int(getattr(owner, count_name)) except Exception: continue getter = getattr(owner, getter_name, None) if not callable(getter): continue for index in range(count): try: yield getter(index) except Exception: continue return def pst_message_folder_path(raw_value: object) -> str | None: if raw_value is None: return None if isinstance(raw_value, (list, tuple)): parts = [normalize_whitespace(str(part)) for part in raw_value if normalize_whitespace(str(part))] return "/".join(parts) or None normalized = normalize_whitespace(str(raw_value).replace("\\", "/")) normalized = re.sub(r"/+", "/", normalized).strip("/") return normalized or None def normalize_pst_chat_thread_id(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")) if not normalized or not normalized.startswith("19:"): return None suffix_index = normalized.lower().find(";messageid=") if suffix_index >= 0: normalized = normalized[:suffix_index] return normalized or None def parse_pst_json_object(raw_value: object) -> dict[str, object] | None: normalized = normalize_whitespace(str(raw_value or "")) if not normalized.startswith("{") or not normalized.endswith("}"): return None try: parsed = json.loads(normalized) except Exception: return None return parsed if isinstance(parsed, dict) else None def normalize_pst_aad_object_id(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")).lower() if not normalized: return None if normalized.startswith(PST_TEAMS_ORGID_PREFIX): normalized = normalized[len(PST_TEAMS_ORGID_PREFIX) :] if re.fullmatch(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", normalized): return normalized if re.fullmatch(r"[0-9a-f]{32}", normalized): return ( f"{normalized[0:8]}-{normalized[8:12]}-{normalized[12:16]}-" f"{normalized[16:20]}-{normalized[20:32]}" ) return None def pst_chat_participant_display_name(payload: object) -> str | None: if not isinstance(payload, dict): return None normalized_name = normalize_participant_token(payload.get("Name")) if normalized_name: return normalized_name normalized_email = normalize_participant_token(payload.get("EmailAddress")) if normalized_email: return normalized_email.lower() return None def pst_chat_participant_entity_hint( payload: object, *, source_kind: str = "pst_teams", ) -> dict[str, object] | None: if not isinstance(payload, dict): return None normalized_name = normalize_entity_text(payload.get("Name") or "") normalized_email = normalize_entity_email(payload.get("EmailAddress")) aad_object_id = normalize_pst_aad_object_id( payload.get("ExternalDirectoryObjectId") or payload.get("AadObjectId") or payload.get("AADObjectId") or payload.get("ObjectId") ) if normalized_name and normalized_email: display_value = f"{normalized_name} <{normalized_email}>" else: display_value = normalized_name or normalized_email or aad_object_id if not display_value: return None identifiers: list[dict[str, object]] = [] if aad_object_id: identifiers.append( { "identifier_type": "external_id", "identifier_name": "aad_oid", "display_value": aad_object_id, "normalized_value": normalize_entity_lookup_text(aad_object_id), "is_verified": 1, "source_kind": source_kind, } ) hint: dict[str, object] = { "display_value": display_value, "source_kind": source_kind, } if identifiers: hint["identifiers"] = identifiers return hint def pst_chat_participant_entity_hint_key(hint: dict[str, object]) -> str: display_key = normalize_entity_lookup_text(hint.get("display_value") or "") for identifier in list(hint.get("identifiers") or []): if isinstance(identifier, dict): identifier_key = entity_candidate_identifier_key(identifier) if identifier_key: return identifier_key return display_key def append_unique_pst_aad_object_id(target: list[str], value: object) -> None: aad_object_id = normalize_pst_aad_object_id(value) if aad_object_id and aad_object_id not in target: target.append(aad_object_id) def append_unique_normalized_email(target: list[str], value: object) -> None: normalized_email = normalize_entity_email(value) if normalized_email and normalized_email not in target: target.append(normalized_email) def append_unique_entity_text(target: list[str], value: object) -> None: normalized_text = normalize_entity_text(value or "") if normalized_text and normalized_text not in target: target.append(normalized_text) def pst_legacy_exchange_dn_aad_object_ids(value: object) -> list[str]: text = normalize_whitespace(str(value or "")) if not text: return [] normalized_text = text.lower() marker = "/cn=recipients/cn=" marker_index = normalized_text.rfind(marker) if marker_index < 0: return [] tail = text[marker_index + len(marker) :] aad_object_ids: list[str] = [] for match in PST_DASHED_GUID_PATTERN.finditer(tail): append_unique_pst_aad_object_id(aad_object_ids, match.group(0)) return aad_object_ids def pst_json_payload_child(payload: dict[str, object], key: str) -> dict[str, object] | None: value = payload.get(key) if isinstance(value, dict): return value return parse_pst_json_object(value) def pst_chat_sender_aad_object_ids_from_payload( payload: object, *, depth: int = 0, ) -> list[str]: if not isinstance(payload, dict): return [] aad_object_ids: list[str] = [] for key in ("SenderId", "senderId", "messageFrom", "MessageFrom"): append_unique_pst_aad_object_id(aad_object_ids, payload.get(key)) from_payload = payload.get("from") or payload.get("From") if isinstance(from_payload, dict): for key in ("internalId", "InternalId", "id", "Id", "userId", "UserId"): append_unique_pst_aad_object_id(aad_object_ids, from_payload.get(key)) if depth < 2: for key in ("Context", "ItemData"): child_payload = pst_json_payload_child(payload, key) for aad_object_id in pst_chat_sender_aad_object_ids_from_payload(child_payload, depth=depth + 1): append_unique_pst_aad_object_id(aad_object_ids, aad_object_id) return aad_object_ids def append_pst_chat_sender_entity_hint( target: list[dict[str, object]], seen_keys: set[str], *, sender_display_name: object, sender_email_address: object, sender_aad_object_ids: list[str], ) -> None: normalized_name = normalize_entity_text(sender_display_name or "") normalized_email = normalize_entity_email(sender_email_address) aad_object_id = sender_aad_object_ids[0] if sender_aad_object_ids else None if not any((normalized_name, normalized_email, aad_object_id)): return hint = pst_chat_participant_entity_hint( { "Name": normalized_name, "EmailAddress": normalized_email, "ExternalDirectoryObjectId": aad_object_id, }, source_kind="pst_teams", ) if hint is None: return hint_key = pst_chat_participant_entity_hint_key(hint) if not hint_key or hint_key in seen_keys: return seen_keys.add(hint_key) target.append(hint) def pst_chat_participant_entity_hints_from_payload(payload: object) -> list[dict[str, object]]: if not isinstance(payload, dict): return [] hints: list[dict[str, object]] = [] sender_hint = pst_chat_participant_entity_hint(payload.get("Sender")) if sender_hint is not None: hints.append(sender_hint) recipients = payload.get("Recipients") if isinstance(recipients, list): for recipient in recipients: recipient_hint = pst_chat_participant_entity_hint(recipient) if recipient_hint is not None: hints.append(recipient_hint) return hints def append_pst_chat_participant_entity_hints( target: list[dict[str, object]], seen_keys: set[str], payload: object, ) -> None: for hint in pst_chat_participant_entity_hints_from_payload(payload): hint_key = pst_chat_participant_entity_hint_key(hint) if not hint_key or hint_key in seen_keys: continue seen_keys.add(hint_key) target.append(hint) def pst_chat_participant_names_from_payload(payload: object) -> list[str]: if not isinstance(payload, dict): return [] participant_names: list[object] = [] sender_display_name = pst_chat_participant_display_name(payload.get("Sender")) if sender_display_name: participant_names.append(sender_display_name) recipients = payload.get("Recipients") if isinstance(recipients, list): for recipient in recipients: recipient_display_name = pst_chat_participant_display_name(recipient) if recipient_display_name: participant_names.append(recipient_display_name) return sorted_unique_display_names(participant_names) def pst_message_mapi_entity_hint( *, names: list[str], emails: list[str], aad_object_ids: list[str], ) -> dict[str, object] | None: normalized_name = normalize_entity_text(names[0] if names else "") normalized_email = normalize_entity_email(emails[0] if emails else None) aad_object_id = aad_object_ids[0] if aad_object_ids else None return pst_chat_participant_entity_hint( { "Name": normalized_name, "EmailAddress": normalized_email, "ExternalDirectoryObjectId": aad_object_id, }, source_kind="pst_mapi", ) def merge_pst_entity_hint( existing_hints: object, *, role: str, hint: dict[str, object] | None, ) -> dict[str, object]: hints = dict(existing_hints) if isinstance(existing_hints, dict) else {} if hint is None: return hints role_hints = [ item for item in list(hints.get(role) or []) if isinstance(item, dict) ] hint_key = pst_chat_participant_entity_hint_key(hint) if hint_key: for item in role_hints: if pst_chat_participant_entity_hint_key(item) == hint_key: return hints role_hints.append(hint) hints[role] = role_hints return hints def merge_pst_entity_hint_payloads( existing_hints: object, incoming_hints: object, *, roles: set[str] | None = None, ) -> dict[str, object]: hints = dict(existing_hints) if isinstance(existing_hints, dict) else {} if not isinstance(incoming_hints, dict): return hints enabled_roles = set(roles) if roles is not None else set(str(role) for role in incoming_hints.keys()) for role in enabled_roles: for hint in list(incoming_hints.get(role) or []): if isinstance(hint, dict): hints = merge_pst_entity_hint(hints, role=role, hint=hint) return hints def pst_message_mapi_entity_hints(message: object) -> dict[str, object]: identity_parts = { "author": { "names": [], "emails": [], "aad_object_ids": [], }, "recipient": { "names": [], "emails": [], "aad_object_ids": [], }, } append_unique_entity_text( identity_parts["author"]["names"], getattr(message, "sender_name", "") or getattr(message, "sender", "") or "", ) append_unique_normalized_email( identity_parts["author"]["emails"], getattr(message, "sender_email_address", "") or "", ) for record_set in pst_record_sets(message): for entry in pst_record_entries(record_set): try: entry_type = int(getattr(entry, "entry_type", 0) or 0) except Exception: entry_type = 0 decoded_value = decode_pst_record_entry_value(entry) if not decoded_value: continue if entry_type in PST_MAPI_AUTHOR_DISPLAY_NAME_ENTRY_TYPES: append_unique_entity_text(identity_parts["author"]["names"], decoded_value) elif entry_type in PST_MAPI_RECIPIENT_DISPLAY_NAME_ENTRY_TYPES: append_unique_entity_text(identity_parts["recipient"]["names"], decoded_value) if entry_type in PST_MAPI_AUTHOR_SMTP_ADDRESS_ENTRY_TYPES: append_unique_normalized_email(identity_parts["author"]["emails"], decoded_value) elif entry_type in PST_MAPI_RECIPIENT_SMTP_ADDRESS_ENTRY_TYPES: append_unique_normalized_email(identity_parts["recipient"]["emails"], decoded_value) if entry_type in PST_MAPI_AUTHOR_LEGACY_DN_ENTRY_TYPES: for aad_object_id in pst_legacy_exchange_dn_aad_object_ids(decoded_value): append_unique_pst_aad_object_id(identity_parts["author"]["aad_object_ids"], aad_object_id) elif entry_type in PST_MAPI_RECIPIENT_LEGACY_DN_ENTRY_TYPES: for aad_object_id in pst_legacy_exchange_dn_aad_object_ids(decoded_value): append_unique_pst_aad_object_id(identity_parts["recipient"]["aad_object_ids"], aad_object_id) author_hint = pst_message_mapi_entity_hint( names=identity_parts["author"]["names"], emails=identity_parts["author"]["emails"], aad_object_ids=identity_parts["author"]["aad_object_ids"], ) recipient_hint = pst_message_mapi_entity_hint( names=identity_parts["recipient"]["names"], emails=identity_parts["recipient"]["emails"], aad_object_ids=identity_parts["recipient"]["aad_object_ids"], ) hints: dict[str, object] = {} if author_hint is not None: hints = merge_pst_entity_hint(hints, role="author", hint=author_hint) hints = merge_pst_entity_hint(hints, role="participants", hint=author_hint) if recipient_hint is not None: hints = merge_pst_entity_hint(hints, role="recipients", hint=recipient_hint) hints = merge_pst_entity_hint(hints, role="participants", hint=recipient_hint) return hints def pst_message_may_have_chat_threading(folder_path: object, message_class: object) -> bool: if pst_folder_path_contains(folder_path, "/skypespacesdata/teamsmeetings/"): return False if pst_folder_path_contains(folder_path, "/teamsmessagesdata/"): return True if pst_folder_path_contains(folder_path, "/conversation history/"): return True normalized_message_class = normalize_whitespace(str(message_class or "")).lower() return "skypeteams" in normalized_message_class or "microsoft.conversation" in normalized_message_class def extract_pst_chat_threading(message: object) -> dict[str, object] | None: thread_id = None message_id = None parent_message_id = None thread_type = None participant_names: list[object] = [] participant_entity_hints: list[dict[str, object]] = [] seen_participant_entity_hint_keys: set[str] = set() sender_display_name = normalize_entity_text( getattr(message, "sender_name", "") or getattr(message, "sender", "") or "" ) sender_email_address = normalize_entity_email(getattr(message, "sender_email_address", "") or "") sender_aad_object_ids: list[str] = [] for record_set in pst_record_sets(message): for entry in pst_record_entries(record_set): decoded_value = decode_pst_record_entry_value(entry) normalized_value = normalize_whitespace(str(decoded_value or "")) if not normalized_value: continue try: entry_type = int(getattr(entry, "entry_type", 0) or 0) except Exception: entry_type = 0 if entry_type in PST_SENDER_DISPLAY_NAME_ENTRY_TYPES and not sender_display_name: sender_display_name = normalize_entity_text(normalized_value) elif entry_type in PST_SENDER_EMAIL_ADDRESS_ENTRY_TYPES and not sender_email_address: sender_email_address = normalize_entity_email(normalized_value) if normalized_value.lower().startswith(PST_TEAMS_ORGID_PREFIX): append_unique_pst_aad_object_id(sender_aad_object_ids, normalized_value) candidate_thread_id = normalize_pst_chat_thread_id(normalized_value) if candidate_thread_id and thread_id is None: thread_id = candidate_thread_id if message_id is None and normalized_value.startswith("19:") and ";messageid=" in normalized_value.lower(): message_id = normalize_pst_identifier(normalized_value.rsplit("=", 1)[-1]) parsed_payload = parse_pst_json_object(normalized_value) if parsed_payload is None: continue parsed_thread_id = normalize_pst_chat_thread_id(parsed_payload.get("ThreadId")) if parsed_thread_id: thread_id = parsed_thread_id parsed_message_id = normalize_pst_identifier(parsed_payload.get("MessageId")) if parsed_message_id: message_id = parsed_message_id parsed_parent_message_id = normalize_pst_identifier(parsed_payload.get("ParentMessageId")) if parsed_parent_message_id: parent_message_id = parsed_parent_message_id normalized_thread_type = normalize_whitespace(str(parsed_payload.get("ThreadType") or "")).lower() if normalized_thread_type: thread_type = normalized_thread_type participant_names.extend(pst_chat_participant_names_from_payload(parsed_payload)) append_pst_chat_participant_entity_hints( participant_entity_hints, seen_participant_entity_hint_keys, parsed_payload, ) for aad_object_id in pst_chat_sender_aad_object_ids_from_payload(parsed_payload): append_unique_pst_aad_object_id(sender_aad_object_ids, aad_object_id) has_threading_signal = any( ( thread_id, message_id, parent_message_id, thread_type, participant_names, participant_entity_hints, sender_aad_object_ids, ) ) if has_threading_signal: if not participant_names: sender_participant_name = normalize_participant_token(sender_display_name) or sender_email_address if sender_participant_name: participant_names.append(sender_participant_name) append_pst_chat_sender_entity_hint( participant_entity_hints, seen_participant_entity_hint_keys, sender_display_name=sender_display_name, sender_email_address=sender_email_address, sender_aad_object_ids=sender_aad_object_ids, ) normalized_participants = sorted_unique_display_names(participant_names) if not any((thread_id, message_id, parent_message_id, thread_type, normalized_participants, participant_entity_hints)): return None payload: dict[str, object] = { "thread_id": thread_id, "message_id": message_id, "parent_message_id": parent_message_id, "thread_type": thread_type, "participants": normalized_participants, } if participant_entity_hints: payload["participant_entity_hints"] = participant_entity_hints return payload def pst_message_author(message: object) -> str | None: sender = normalize_whitespace(str(getattr(message, "sender_name", "") or getattr(message, "sender", "") or "")) sender_email = normalize_whitespace(str(getattr(message, "sender_email_address", "") or "")) if sender and sender_email: return f"{sender} <{sender_email}>" return sender or sender_email or None def pst_message_recipients(message: object, transport_headers: object = None) -> str | None: parts = [ normalize_whitespace(str(getattr(message, "display_to", "") or "")), normalize_whitespace(str(getattr(message, "display_cc", "") or "")), normalize_whitespace(str(getattr(message, "display_bcc", "") or "")), ] recipients = ", ".join(part for part in parts if part) if recipients: return recipients return extract_email_recipients_from_headers(transport_headers) def pst_attachment_record_entry_text(attachment: object, *entry_types: int) -> str | None: ordered_types = [int(entry_type) for entry_type in entry_types] if not ordered_types: return None found_values: dict[int, str] = {} for record_set in pst_record_sets(attachment): for entry in pst_record_entries(record_set): try: entry_type = int(getattr(entry, "entry_type", 0) or 0) except Exception: continue if entry_type not in ordered_types or entry_type in found_values: continue decoded = decode_pst_record_entry_value(entry) if decoded: found_values[entry_type] = decoded for entry_type in ordered_types: if entry_type in found_values: return found_values[entry_type] return None def pst_attachment_declared_file_name(attachment: object) -> str | None: raw_name = None if isinstance(attachment, dict): raw_name = attachment.get("file_name") or attachment.get("name") or attachment.get("filename") else: for attr_name in ("long_filename", "filename", "name", "display_name"): value = getattr(attachment, attr_name, None) if value: raw_name = value break if raw_name is None: raw_name = pst_attachment_record_entry_text( attachment, PST_PROP_ATTACH_LONG_FILENAME, PST_PROP_ATTACH_FILENAME, PST_PROP_DISPLAY_NAME, ) normalized = normalize_whitespace(raw_name if isinstance(raw_name, str) else "") if normalized: return normalized return None def pst_attachment_extension(attachment: object) -> str | None: if isinstance(attachment, dict): for key in ("preferred_extension", "extension", "file_type"): normalized = normalize_file_type_name(attachment.get(key)) if normalized: return normalized return None for attr_name in ("extension", "attach_extension"): normalized = normalize_file_type_name(getattr(attachment, attr_name, None)) if normalized: return normalized return normalize_file_type_name(pst_attachment_record_entry_text(attachment, PST_PROP_ATTACH_EXTENSION)) def pst_attachment_content_type(attachment: object) -> str | None: if isinstance(attachment, dict): return normalize_mime_type(attachment.get("content_type") or attachment.get("mime_type")) for attr_name in ("mime_tag", "attach_mime_tag", "mime_type", "content_type"): normalized = normalize_mime_type(getattr(attachment, attr_name, None)) if normalized: return normalized return normalize_mime_type(pst_attachment_record_entry_text(attachment, PST_PROP_ATTACH_MIME_TAG)) def pst_attachment_file_name( attachment: object, ordinal: int, *, payload: bytes | None = None, ) -> str: raw_name = pst_attachment_declared_file_name(attachment) return normalize_attachment_filename( raw_name, ordinal, payload=payload, content_type=pst_attachment_content_type(attachment), preferred_extension=pst_attachment_extension(attachment), ) def pst_message_html_body(message: object) -> str | None: for attr_name in ("html_body", "htmlBody"): value = getattr(message, attr_name, None) if value is None: continue if isinstance(value, bytes): return value.decode("utf-8", errors="replace") return str(value) return None def pst_message_text_body(message: object, html_body: str | None) -> str | None: for attr_name in ("plain_text_body", "body", "text_body"): value = getattr(message, attr_name, None) if value is None: continue if isinstance(value, bytes): decoded, _, _ = decode_bytes(value) text = normalize_whitespace(decoded) else: text = normalize_whitespace(str(value)) if text: return text if html_body: return strip_html_tags(html_body) return None def pst_message_transport_headers(message: object) -> str | None: for attr_name in ("transport_headers",): value = getattr(message, attr_name, None) if value is None: continue if isinstance(value, bytes): return value.decode("utf-8", errors="replace") normalized = str(value) if normalize_whitespace(normalized): return normalized for method_name in ("get_transport_headers",): method = getattr(message, method_name, None) if not callable(method): continue try: value = method() except Exception: continue if value is None: continue if isinstance(value, bytes): return value.decode("utf-8", errors="replace") normalized = str(value) if normalize_whitespace(normalized): return normalized return None def pst_message_conversation_topic(message: object) -> str | None: for attr_name in ("conversation_topic",): value = getattr(message, attr_name, None) normalized = normalize_whitespace(str(value or "")) if normalized: return normalized for method_name in ("get_conversation_topic",): method = getattr(message, method_name, None) if not callable(method): continue try: value = method() except Exception: continue normalized = normalize_whitespace(str(value or "")) if normalized: return normalized return None def pst_record_sets(owner: object): yield from iter_pst_collection( owner, list_attrs=("record_sets",), count_getter_pairs=(("number_of_record_sets", "get_record_set"),), ) def pst_record_entries(record_set: object): yield from iter_pst_collection( record_set, list_attrs=("entries", "record_entries"), count_getter_pairs=(("number_of_entries", "get_entry"), ("number_of_record_entries", "get_record_entry")), ) def pst_debug_value_preview(raw_value: object, *, max_chars: int = PST_DEBUG_MAX_TEXT_VALUE_CHARS) -> str | None: normalized = normalize_whitespace(str(raw_value or "")) if not normalized: return None if len(normalized) <= max_chars: return normalized return normalized[: max_chars - 3].rstrip() + "..." def pst_debug_record_entry_payload(entry: object, *, entry_index: int) -> dict[str, object]: payload: dict[str, object] = {"entry_index": int(entry_index)} for field_name in ( "identifier", "name", "entry_name", "property_name", "named_property_name", "guid", "property_guid", "property_set_guid", "named_property_guid", ): normalized = pst_debug_value_preview(getattr(entry, field_name, None)) if normalized: payload[field_name] = normalized try: entry_type = int(getattr(entry, "entry_type", 0) or 0) except Exception: entry_type = None if entry_type is not None: payload["entry_type"] = entry_type payload["entry_type_hex"] = f"0x{entry_type:04X}" try: value_type = int(getattr(entry, "value_type", 0) or 0) except Exception: value_type = None if value_type is not None: payload["value_type"] = value_type payload["value_type_hex"] = f"0x{value_type:04X}" data = getattr(entry, "data", None) if isinstance(data, (bytes, bytearray)): payload["byte_length"] = len(data) payload["hex_preview"] = bytes(data[:PST_DEBUG_MAX_HEX_PREVIEW_BYTES]).hex() decoded_value = decode_pst_record_entry_value(entry) if decoded_value: payload["decoded_value"] = pst_debug_value_preview(decoded_value) return payload def pst_debug_entry_scope_candidate(payload: dict[str, object]) -> str | None: decoded_value = normalize_whitespace(str(payload.get("decoded_value") or "")) if not decoded_value or len(decoded_value) > PST_DEBUG_MAX_TEXT_VALUE_CHARS: return None if decoded_value.lower().startswith("19:"): return decoded_value if PST_DEBUG_GUID_VALUE_PATTERN.search(decoded_value): return decoded_value name_blob = " ".join( normalize_whitespace(str(payload.get(field_name) or "")) for field_name in ( "name", "entry_name", "property_name", "named_property_name", ) ) if PST_DEBUG_SCOPE_NAME_PATTERN.search(name_blob): return decoded_value return None def pst_debug_interesting_entry(payload: dict[str, object]) -> bool: candidate = pst_debug_entry_scope_candidate(payload) if candidate: return True name_blob = " ".join( normalize_whitespace(str(payload.get(field_name) or "")) for field_name in ( "name", "entry_name", "property_name", "named_property_name", ) ) if PST_DEBUG_SCOPE_NAME_PATTERN.search(name_blob): return True decoded_value = normalize_whitespace(str(payload.get("decoded_value") or "")) return bool(decoded_value and decoded_value.lower().startswith("https://teams.microsoft.com/")) def pst_debug_record_sets_payloads( owner: object, *, max_record_entries: int = 128, ) -> tuple[list[dict[str, object]], list[dict[str, object]], list[str]]: record_set_payloads: list[dict[str, object]] = [] interesting_entries: list[dict[str, object]] = [] candidate_scope_values: list[str] = [] seen_candidates: set[str] = set() for record_set_index, record_set in enumerate(pst_record_sets(owner), start=1): entry_payloads: list[dict[str, object]] = [] entry_total = 0 truncated = False for entry_index, entry in enumerate(pst_record_entries(record_set), start=1): entry_total += 1 if entry_index > max_record_entries: truncated = True continue entry_payload = pst_debug_record_entry_payload(entry, entry_index=entry_index) entry_payloads.append(entry_payload) if pst_debug_interesting_entry(entry_payload): interesting_entry = { "record_set_index": int(record_set_index), **entry_payload, } interesting_entries.append(interesting_entry) candidate = pst_debug_entry_scope_candidate(entry_payload) if candidate and candidate not in seen_candidates: seen_candidates.add(candidate) candidate_scope_values.append(candidate) record_set_payloads.append( { "record_set_index": int(record_set_index), "entry_count": int(entry_total), "entries_truncated": truncated, "entries": entry_payloads, } ) return record_set_payloads, interesting_entries, candidate_scope_values def iter_pst_raw_messages( path: Path, *, include_debug_record_sets: bool = False, max_record_entries: int = 128, ): pypff_module = dependency_guard("pypff", "libpff-python", "pst") def _iter_folder(folder: object, ancestors: list[str]): folder_name = normalize_whitespace(str(getattr(folder, "name", "") or "")) current_ancestors = [*ancestors] if folder_name: current_ancestors.append(folder_name) folder_path = pst_message_folder_path(current_ancestors) for message in iter_pst_collection( folder, list_attrs=("sub_messages", "messages"), count_getter_pairs=(("number_of_sub_messages", "get_sub_message"), ("number_of_messages", "get_message")), ): source_item_id = None for attr_name in ("entry_identifier", "entry_identifier_string", "record_key", "search_key", "identifier"): source_item_id = normalize_pst_identifier(getattr(message, attr_name, None)) if source_item_id: break if not source_item_id: raise RetrieverError(f"PST message is missing a stable item identifier in {path}") html_body = pst_message_html_body(message) attachments: list[dict[str, object]] = [] for ordinal, attachment in enumerate( iter_pst_collection( message, list_attrs=("attachments",), count_getter_pairs=(("number_of_attachments", "get_attachment"),), ), start=1, ): payload = coerce_pst_attachment_payload(attachment) if payload is None: continue content_type = pst_attachment_content_type(attachment) attachments.append( { "file_name": pst_attachment_file_name(attachment, ordinal, payload=payload), "ordinal": ordinal, "payload": payload, "file_hash": sha256_bytes(payload), "content_type": content_type, "content_id": pst_attachment_content_id(attachment), } ) message_class = normalize_whitespace(str(getattr(message, "message_class", "") or "")) or None entity_hints = pst_message_mapi_entity_hints(message) chat_threading = ( extract_pst_chat_threading(message) if pst_message_may_have_chat_threading(folder_path, message_class) else None ) debug_record_sets: list[dict[str, object]] = [] debug_interesting_entries: list[dict[str, object]] = [] debug_candidate_scope_values: list[str] = [] if include_debug_record_sets: ( debug_record_sets, debug_interesting_entries, debug_candidate_scope_values, ) = pst_debug_record_sets_payloads( message, max_record_entries=max_record_entries, ) transport_headers = pst_message_transport_headers(message) yield { "source_item_id": source_item_id, "folder_path": folder_path, "message_class": message_class, "subject": normalize_whitespace(str(getattr(message, "subject", "") or "")) or None, "conversation_topic": pst_message_conversation_topic(message), "transport_headers": transport_headers, "author": pst_message_author(message), "recipients": pst_message_recipients(message, transport_headers), "entity_hints": entity_hints, "date_created": normalize_datetime( getattr(message, "delivery_time", None) or getattr(message, "client_submit_time", None) or getattr(message, "creation_time", None) ), "text_body": pst_message_text_body(message, html_body), "html_body": html_body, "attachments": attachments, "chat_threading": chat_threading, "high_level_identifiers": { attr_name: normalize_pst_identifier(getattr(message, attr_name, None)) for attr_name in ("entry_identifier", "entry_identifier_string", "record_key", "search_key", "identifier") if normalize_pst_identifier(getattr(message, attr_name, None)) }, "debug_record_sets": debug_record_sets, "debug_interesting_entries": debug_interesting_entries, "debug_candidate_scope_values": debug_candidate_scope_values, } for child_folder in iter_pst_collection( folder, list_attrs=("sub_folders", "folders"), count_getter_pairs=(("number_of_sub_folders", "get_sub_folder"), ("number_of_folders", "get_folder")), ): yield from _iter_folder(child_folder, current_ancestors) pst_file = pypff_module.file() # type: ignore[union-attr] pst_file.open(str(path)) try: root_folder = pst_file.get_root_folder() yield from _iter_folder(root_folder, []) finally: try: pst_file.close() except Exception: pass def iter_pst_debug_messages( path: Path, *, max_record_entries: int = 128, ): for message_dict in iter_pst_raw_messages( path, include_debug_record_sets=True, max_record_entries=max_record_entries, ): chat_text = str(message_dict.get("text_body") or "") or ( strip_html_tags(str(message_dict.get("html_body") or "")) if message_dict.get("html_body") else "" ) chat_metadata = extract_chat_transcript_metadata(chat_text) message_kind = classify_pst_message_kind(message_dict, chat_metadata) payload = { "source_item_id": message_dict["source_item_id"], "folder_path": message_dict.get("folder_path"), "message_kind": message_kind, "message_class": message_dict.get("message_class"), "subject": message_dict.get("subject"), "conversation_topic": message_dict.get("conversation_topic"), "author": message_dict.get("author"), "recipients": message_dict.get("recipients"), "date_created": message_dict.get("date_created"), "high_level_identifiers": dict(message_dict.get("high_level_identifiers") or {}), "candidate_scope_values": list(message_dict.get("debug_candidate_scope_values") or []), "interesting_properties": list(message_dict.get("debug_interesting_entries") or []), "record_sets": list(message_dict.get("debug_record_sets") or []), } if chat_metadata: payload["chat_metadata"] = { key: chat_metadata[key] for key in ("title", "participants", "date_created", "date_modified", "message_count") if key in chat_metadata and chat_metadata.get(key) not in (None, "") } yield payload def decode_pst_record_entry_value(entry: object) -> str | None: data = getattr(entry, "data", None) if isinstance(data, (bytes, bytearray)): value_type = None try: value_type = int(getattr(entry, "value_type", 0) or 0) except Exception: value_type = None raw = bytes(data) try: if value_type == 0x001F: decoded = raw.decode("utf-16-le", errors="replace") elif value_type == 0x001E: decoded = raw.decode("utf-8", errors="replace") else: try: decoded = raw.decode("utf-16-le") except UnicodeDecodeError: decoded = raw.decode("utf-8", errors="replace") except Exception: decoded = raw.decode("utf-8", errors="replace") return decoded.rstrip("\x00").strip() or None for method_name in ("get_data_as_string",): method = getattr(entry, method_name, None) if callable(method): try: value = method() except Exception: continue if value: text = value.decode("utf-8", errors="replace") if isinstance(value, (bytes, bytearray)) else str(value) text = text.rstrip("\x00").strip() if text: return text if data is not None: text = str(data).rstrip("\x00").strip() return text or None return None def pst_attachment_content_id(attachment: object) -> str | None: for attr_name in ("content_id", "attach_content_id", "long_content_id"): direct = getattr(attachment, attr_name, None) normalized = normalize_content_id(direct) if normalized: return normalized record_sets = getattr(attachment, "record_sets", None) if record_sets is None: return None try: record_sets_iter = list(record_sets) except Exception: return None for record_set in record_sets_iter: for entries_attr in ("entries", "record_entries"): entries = getattr(record_set, entries_attr, None) if entries is None: continue try: entries_iter = list(entries) except Exception: continue for entry in entries_iter: try: entry_type = int(getattr(entry, "entry_type", 0) or 0) except Exception: continue if entry_type != PST_PROP_ATTACH_CONTENT_ID: continue decoded = decode_pst_record_entry_value(entry) normalized = normalize_content_id(decoded) if normalized: return normalized break return None def iter_pst_messages(path: Path): for payload in iter_pst_raw_messages(path): yield { "source_item_id": payload["source_item_id"], "folder_path": payload.get("folder_path"), "message_class": payload.get("message_class"), "subject": payload.get("subject"), "conversation_topic": payload.get("conversation_topic"), "transport_headers": payload.get("transport_headers"), "author": payload.get("author"), "recipients": payload.get("recipients"), "entity_hints": payload.get("entity_hints"), "date_created": payload.get("date_created"), "text_body": payload.get("text_body"), "html_body": payload.get("html_body"), "chat_threading": payload.get("chat_threading"), "attachments": list(payload.get("attachments") or []), } def normalize_pst_message(source_rel_path: str, message_dict: dict[str, object]) -> dict[str, object] | None: source_item_id = normalize_source_item_id( message_dict.get("source_item_id") or message_dict.get("entry_identifier") or message_dict.get("identifier") ) normalized_attachments: list[dict[str, object]] = [] for ordinal, raw_attachment in enumerate(list(message_dict.get("attachments") or []), start=1): raw_name = None raw_content_type: object = None raw_extension: object = None raw_content_id: object = None if isinstance(raw_attachment, dict): raw_name = raw_attachment.get("file_name") or raw_attachment.get("name") or raw_attachment.get("filename") raw_content_type = raw_attachment.get("content_type") or raw_attachment.get("mime_type") raw_extension = raw_attachment.get("preferred_extension") or raw_attachment.get("extension") or raw_attachment.get("file_type") raw_content_id = raw_attachment.get("content_id") else: raw_name = pst_attachment_declared_file_name(raw_attachment) raw_content_type = pst_attachment_content_type(raw_attachment) raw_extension = pst_attachment_extension(raw_attachment) raw_content_id = pst_attachment_content_id(raw_attachment) payload = coerce_pst_attachment_payload(raw_attachment) if payload is None: continue attachment_ordinal = ordinal if isinstance(raw_attachment, dict): try: attachment_ordinal = int(raw_attachment.get("ordinal", ordinal)) except Exception: attachment_ordinal = ordinal normalized_attachments.append( { "file_name": normalize_attachment_filename( raw_name if isinstance(raw_name, str) else None, attachment_ordinal, payload=payload, content_type=raw_content_type, preferred_extension=raw_extension, ), "ordinal": attachment_ordinal, "payload": payload, "file_hash": sha256_bytes(payload), "content_type": normalize_mime_type(raw_content_type), "content_id": normalize_content_id(raw_content_id), } ) html_body = message_dict.get("html_body") normalized_subject = normalize_whitespace(str(message_dict.get("subject") or "")) or None normalized_author = normalize_whitespace(str(message_dict.get("author") or "")) or None normalized_recipients = normalize_whitespace(str(message_dict.get("recipients") or "")) or None if normalized_recipients is None: normalized_recipients = extract_email_recipients_from_headers(message_dict.get("transport_headers")) normalized_date_created = normalize_datetime(message_dict.get("date_created")) normalized_text_body = None if message_dict.get("text_body") is None else str(message_dict.get("text_body") or "") normalized_html_body = None if html_body is None else str(html_body) chat_text = normalized_text_body or (strip_html_tags(normalized_html_body) if normalized_html_body else "") chat_metadata = extract_chat_transcript_metadata(chat_text) message_kind = classify_pst_message_kind(message_dict, chat_metadata) if message_kind == "skip": return None if message_kind == "chat": raw_chat_threading = message_dict.get("chat_threading") chat_threading = dict(raw_chat_threading) if isinstance(raw_chat_threading, dict) else {} chat_thread_participants = sorted_unique_display_names(normalize_string_list(chat_threading.get("participants"))) participant_entity_hints = [ dict(item) for item in list(chat_threading.get("participant_entity_hints") or []) if isinstance(item, dict) ] preferred_participants = render_display_name_list(chat_thread_participants) preferred_title = None if normalize_whitespace(str(chat_threading.get("thread_type") or "")).lower() == "chat": preferred_title = render_display_name_title(chat_thread_participants, max_names=4) chat_entries = synthesize_pst_chat_entries( author=normalized_author, date_created=normalized_date_created, text_body=chat_text, chat_metadata=chat_metadata, ) resolved_chat_metadata = synthesize_pst_chat_metadata( subject=normalized_subject, author=normalized_author, date_created=normalized_date_created, text_body=chat_text, chat_metadata=chat_metadata, chat_entries=chat_entries, preferred_title=preferred_title, preferred_participants=preferred_participants, ) extracted = build_chat_extracted_payload( title=preferred_title or normalized_subject, author=normalized_author, date_created=normalized_date_created, text_body=normalized_text_body, html_body=normalized_html_body, attachments=normalized_attachments, preview_file_name=pst_preview_file_name(source_item_id), chat_metadata=resolved_chat_metadata, chat_entries=chat_entries, chat_threading=chat_threading, ) if participant_entity_hints: extracted["entity_hints"] = {"participants": participant_entity_hints} elif message_kind == "calendar": extracted = build_calendar_extracted_payload( subject=normalized_subject, author=normalized_author, recipients=normalized_recipients, date_created=normalized_date_created, text_body=normalized_text_body, html_body=normalized_html_body, attachments=normalized_attachments, preview_file_name=pst_preview_file_name(source_item_id), ) else: email_threading = extract_transport_header_threading( message_dict.get("transport_headers"), subject=normalized_subject, conversation_topic=message_dict.get("conversation_topic"), ) extracted = build_email_extracted_payload( subject=normalized_subject, author=normalized_author, recipients=normalized_recipients, date_created=normalized_date_created, text_body=normalized_text_body, html_body=normalized_html_body, attachments=normalized_attachments, preview_file_name=pst_preview_file_name(source_item_id), email_threading=email_threading, ) mapi_hint_roles: set[str] = set() if normalize_whitespace(str(extracted.get("author") or "")): mapi_hint_roles.add("author") if normalize_whitespace(str(extracted.get("participants") or "")): mapi_hint_roles.add("participants") if normalize_whitespace(str(extracted.get("recipients") or "")): mapi_hint_roles.add("recipients") if mapi_hint_roles: merged_entity_hints = merge_pst_entity_hint_payloads( extracted.get("entity_hints"), message_dict.get("entity_hints"), roles=mapi_hint_roles, ) if merged_entity_hints: extracted["entity_hints"] = merged_entity_hints return { "rel_path": pst_message_rel_path(source_rel_path, source_item_id), "file_name": pst_message_file_name(source_item_id), "file_hash": container_message_payload_hash( { "subject": extracted.get("subject"), "author": extracted.get("author"), "recipients": extracted.get("recipients"), "date_created": extracted.get("date_created"), "text_body": extracted.get("text_content"), "html_body": html_body, "attachments": normalized_attachments, } ), "source_rel_path": source_rel_path, "source_item_id": source_item_id, "source_folder_path": pst_message_folder_path(message_dict.get("folder_path")), "extracted": extracted, } def extract_document(path: Path, include_attachments: bool = True) -> dict[str, object]: file_type = normalize_extension(path) if file_type not in SUPPORTED_FILE_TYPES: raise RetrieverError(f"Unsupported file type: .{file_type or '(none)'}") if file_type in {"csv", "tsv"}: return extract_csv_file(path) if file_type in TEXT_FILE_TYPES: return extract_plain_text_file(path) if file_type in IMAGE_NATIVE_PREVIEW_FILE_TYPES: return extract_native_preview_only_file(path, explicit_content_type="Image") if file_type == "rtf": return extract_rtf_file(path) if file_type == "pdf": return extract_pdf_file(path) if file_type == "docx": return extract_docx_file(path) if file_type == "pptx": return extract_pptx_file(path) if file_type == "xls": return extract_xls_file(path) if file_type == "xlsx": return extract_xlsx_file(path) if file_type == "eml": return extract_eml_file(path, include_attachments=include_attachments) if file_type == "msg": return extract_msg_file(path, include_attachments=include_attachments) if file_type == "mbox": raise RetrieverError("MBOX sources must be ingested through the container ingest pipeline.") if file_type == "pst": raise RetrieverError("PST sources must be ingested through the container ingest pipeline.") raise RetrieverError(f"Unsupported file type: .{file_type}") SLACK_EXPORT_DAY_FILE_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}\.json$") SLACK_EXPORT_CONVERSATION_INDEX_FILES = { "channels.json": "public_channel", "groups.json": "private_channel", "dms.json": "dm", "mpims.json": "mpim", } SLACK_EXPORT_AUXILIARY_FILES = { "users.json", "canvases.json", "integration_logs.json", } SLACK_EXPORT_METADATA_FILES = set(SLACK_EXPORT_CONVERSATION_INDEX_FILES) | SLACK_EXPORT_AUXILIARY_FILES def is_slack_export_day_file(path: Path) -> bool: return path.is_file() and SLACK_EXPORT_DAY_FILE_PATTERN.fullmatch(path.name) is not None def iter_slack_export_day_files(export_root: Path) -> list[Path]: day_files: list[Path] = [] try: children = sorted(export_root.iterdir()) except OSError: return [] for child in children: if not child.is_dir() or child.name == ".retriever": continue try: for item in sorted(child.iterdir()): if is_slack_export_day_file(item): day_files.append(item) except OSError: continue return day_files def slack_export_day_file_sample_valid(path: Path) -> bool: try: decoded, _, _ = decode_bytes(path.read_bytes()) except OSError: return False return extract_slack_chat_json_payload(path, decoded) is not None def detect_slack_export_root(candidate_root: Path) -> dict[str, object] | None: if not candidate_root.is_dir() or ".retriever" in candidate_root.parts: return None if not (candidate_root / "users.json").exists(): return None index_files = [ file_name for file_name in sorted(SLACK_EXPORT_CONVERSATION_INDEX_FILES) if (candidate_root / file_name).exists() ] if not index_files: return None day_files = iter_slack_export_day_files(candidate_root) if not day_files: return None sample_valid = any(slack_export_day_file_sample_valid(path) for path in day_files[: min(3, len(day_files))]) if not sample_valid: return None return { "root": candidate_root, "index_files": index_files, "day_files": day_files, } def find_slack_export_roots( root: Path, recursive: bool, allowed_file_types: set[str] | None, ) -> list[dict[str, object]]: if allowed_file_types is not None and "json" not in allowed_file_types: return [] candidates: set[Path] = set() if recursive: try: for users_path in root.rglob("users.json"): if ".retriever" in users_path.parts: continue candidates.add(users_path.parent) except OSError: return [] else: if (root / "users.json").exists(): candidates.add(root) descriptors: list[dict[str, object]] = [] accepted_roots: list[Path] = [] for candidate in sorted(candidates, key=lambda path: (len(path.parts), path.as_posix())): if any(parent == candidate or parent in candidate.parents for parent in accepted_roots): continue descriptor = detect_slack_export_root(candidate) if descriptor is None: continue descriptors.append(descriptor) accepted_roots.append(candidate) return descriptors def slack_conversation_display_name(folder_name: str, conversation_type: str) -> str: normalized = normalize_whitespace(folder_name) if conversation_type in {"public_channel", "private_channel"} and normalized: return normalized if normalized.startswith("#") else f"#{normalized}" return normalized or "Slack conversation" def load_slack_export_conversation_directory(export_root: Path) -> dict[str, dict[str, str]]: directory: dict[str, dict[str, str]] = {} for file_name, conversation_type in SLACK_EXPORT_CONVERSATION_INDEX_FILES.items(): path = export_root / file_name if not path.exists(): continue try: raw_items = json.loads(path.read_text(encoding="utf-8")) except (OSError, ValueError, json.JSONDecodeError): continue if not isinstance(raw_items, list): continue for item in raw_items: if not isinstance(item, dict): continue folder_name = choose_slack_text(item.get("name"), item.get("id")) if not folder_name: continue folder_key = normalize_whitespace(folder_name).lower() if not folder_key: continue conversation_key = choose_slack_text(item.get("id"), folder_name) or folder_name display_name = slack_conversation_display_name(folder_name, conversation_type) directory[folder_key] = { "conversation_key": conversation_key, "conversation_type": conversation_type, "display_name": display_name, } return directory def slack_conversation_metadata_for_folder( folder_name: str, directory: dict[str, dict[str, str]], ) -> dict[str, str]: folder_key = normalize_whitespace(folder_name).lower() if folder_key in directory: return directory[folder_key] return { "conversation_key": normalize_whitespace(folder_name) or folder_name, "conversation_type": "channel", "display_name": slack_conversation_display_name(folder_name, "channel"), } def slack_reply_thread_rel_path(conversation_key: str, thread_ts: str) -> str: conversation_segment = re.sub(r"[^A-Za-z0-9._-]+", "_", normalize_whitespace(conversation_key) or "conversation") thread_segment = re.sub(r"[^A-Za-z0-9._-]+", "_", normalize_whitespace(thread_ts) or "thread") return ( Path("_retriever") / "logical" / "slack" / conversation_segment / "threads" / f"{thread_segment}.slackthread" ).as_posix() def slack_day_document_title(display_name: str, day_file: Path) -> str: day_token = normalize_whitespace(day_file.stem) try: day_label = date.fromisoformat(day_token).strftime("%b %d, %Y").replace(" 0", " ") except ValueError: day_label = day_token normalized_display_name = normalize_whitespace(display_name) if normalized_display_name and day_label: return f"{normalized_display_name} - {day_label}" return normalized_display_name or day_label or "Slack conversation" def slack_reply_thread_title(display_name: str, root_timestamp: str | None) -> str: timestamp_label = format_chat_preview_timestamp(root_timestamp) normalized_display_name = normalize_whitespace(display_name) or "Slack conversation" if timestamp_label: return f"{normalized_display_name} - thread from {timestamp_label}" return f"{normalized_display_name} - thread" def slack_timestamp_sort_key(value: object) -> tuple[float, str]: raw = normalize_whitespace(str(value or "")) try: return (float(raw), raw) except (TypeError, ValueError): return (float("inf"), raw) def slack_message_has_thread(raw_message: dict[str, object], ts_raw: str, thread_ts_raw: str | None) -> bool: if thread_ts_raw and thread_ts_raw == ts_raw: return True for key in ("reply_count", "reply_users_count"): try: if int(raw_message.get(key) or 0) > 0: return True except (TypeError, ValueError): continue if raw_message.get("latest_reply") is not None: return True replies = raw_message.get("replies") return isinstance(replies, list) and bool(replies) def parse_slack_int(value: object) -> int: normalized = normalize_whitespace(str(value or "")) if not normalized: return 0 try: return int(normalized) except ValueError: return 0 def slack_actor_entity_hint( actor_info: dict[str, str | None], *, identifier_scope: str | None, ) -> dict[str, object] | None: slack_user_id = normalize_whitespace(str(actor_info.get("slack_user_id") or "")) speaker_name = normalize_entity_text(actor_info.get("speaker_name") or "") if not slack_user_id or not speaker_name: return None identifiers: list[dict[str, object]] = [] external_identifier: dict[str, object] = { "identifier_type": "external_id", "identifier_name": "slack_user_id", "display_value": slack_user_id, "normalized_value": normalize_entity_lookup_text(slack_user_id), "is_verified": 1, } normalized_scope = normalize_entity_lookup_text(identifier_scope or "") if normalized_scope: external_identifier["identifier_scope"] = normalized_scope identifiers.append(external_identifier) email = normalize_entity_email(actor_info.get("email") or "") if email: identifiers.append( { "identifier_type": "email", "display_value": email, "normalized_value": email, "is_verified": 1, } ) handle = normalize_entity_handle(actor_info.get("handle") or "") if handle and normalized_scope: identifiers.append( { "identifier_type": "handle", "display_value": f"@{handle}", "normalized_value": handle, "provider": "slack", "provider_scope": normalized_scope, "is_verified": 1, } ) return { "display_value": speaker_name, "identifiers": identifiers, } def load_slack_day_messages( day_file: Path, *, rel_path: str, user_directory: dict[str, dict[str, str | None]], entity_identifier_scope: str | None = None, ) -> list[dict[str, object]]: try: raw_value = json.loads(day_file.read_text(encoding="utf-8")) except (OSError, ValueError, json.JSONDecodeError): return [] candidate_items = raw_value.get("messages") if isinstance(raw_value, dict) else raw_value if not isinstance(candidate_items, list): return [] messages: list[dict[str, object]] = [] for ordinal, item in enumerate(candidate_items): if not isinstance(item, dict): continue if normalize_whitespace(str(item.get("type") or "")).lower() != "message": continue ts_raw = choose_slack_text(item.get("ts")) if not ts_raw: continue body = render_slack_text(choose_slack_text(item.get("text")) or "", user_directory) if not body: continue actor_info = slack_message_actor_info(item, user_directory) speaker = actor_info.get("speaker_name") or "Slack message" entity_hint = slack_actor_entity_hint( actor_info, identifier_scope=entity_identifier_scope, ) timestamp = normalize_slack_timestamp(ts_raw) thread_ts_raw = choose_slack_text(item.get("thread_ts")) is_reply = bool(thread_ts_raw and thread_ts_raw != ts_raw) messages.append( { "avatar_color": actor_info.get("avatar_color"), "speaker": speaker, "body": body, "timestamp": timestamp, "timestamp_label": format_chat_preview_timestamp(timestamp), "avatar_label": chat_avatar_initials(speaker), "ts": ts_raw, "thread_ts": thread_ts_raw, "reply_count": parse_slack_int(item.get("reply_count")), "is_reply": is_reply, "is_thread_root": (not is_reply) and slack_message_has_thread(item, ts_raw, thread_ts_raw), "day_rel_path": rel_path, "day_file_name": day_file.name, "ordinal": ordinal, "entity_hint": entity_hint, } ) return sorted( messages, key=lambda message: ( slack_timestamp_sort_key(message.get("ts")), int(message.get("ordinal") or 0), ), ) def build_slack_transcript_components(messages: list[dict[str, object]]) -> dict[str, object]: participants: list[str] = [] seen_participants: set[str] = set() participant_entity_hints: list[dict[str, object]] = [] seen_participant_hint_keys: set[str] = set() transcript_lines: list[str] = [] chat_entries: list[dict[str, object]] = [] first_timestamp: str | None = None last_timestamp: str | None = None timestamped_message_count = 0 for message in sorted(messages, key=lambda item: (slack_timestamp_sort_key(item.get("ts")), int(item.get("ordinal") or 0))): speaker = normalize_whitespace(str(message.get("speaker") or "")) or "Slack message" normalized_speaker = speaker.lower() if normalized_speaker not in seen_participants: seen_participants.add(normalized_speaker) participants.append(speaker) entity_hint = message.get("entity_hint") if isinstance(entity_hint, dict): hint_key = normalize_entity_lookup_text(entity_hint.get("display_value") or "") for identifier in list(entity_hint.get("identifiers") or []): if isinstance(identifier, dict) and identifier.get("identifier_type") == "external_id": hint_key = "|".join( [ "external_id", normalize_entity_lookup_text(identifier.get("identifier_name") or ""), normalize_entity_lookup_text(identifier.get("identifier_scope") or ""), normalize_entity_lookup_text(identifier.get("normalized_value") or identifier.get("display_value") or ""), ] ) break if hint_key and hint_key not in seen_participant_hint_keys: seen_participant_hint_keys.add(hint_key) participant_entity_hints.append(entity_hint) body = normalize_whitespace(str(message.get("body") or "")) timestamp = normalize_whitespace(str(message.get("timestamp") or "")) or None if timestamp: timestamped_message_count += 1 if first_timestamp is None: first_timestamp = timestamp last_timestamp = timestamp transcript_lines.append(f"[{timestamp}] {speaker}: {body}") else: transcript_lines.append(f"{speaker}: {body}") chat_entries.append( { "avatar_color": message.get("avatar_color"), "speaker": speaker, "body": body, "timestamp": timestamp, "timestamp_label": message.get("timestamp_label"), "avatar_label": message.get("avatar_label"), } ) return { "participants": ", ".join(participants) or None, "text_content": normalize_whitespace("\n".join(transcript_lines)), "chat_entries": chat_entries, "date_created": first_timestamp, "date_modified": last_timestamp if last_timestamp and last_timestamp != first_timestamp else None, "message_count": len(chat_entries), "timestamped_message_count": timestamped_message_count, "entity_hints": {"participants": participant_entity_hints} if participant_entity_hints else {}, } def slack_root_message_key(conversation_key: str, thread_ts: str) -> str: return f"{normalize_whitespace(conversation_key) or conversation_key}:{normalize_whitespace(thread_ts) or thread_ts}" def slack_preview_file_name(file_name: str) -> str: return f"{Path(file_name).stem}.html" def build_slack_chat_payload( *, title: str, preview_file_name: str, messages: list[dict[str, object]], placeholder_text: str | None = None, ) -> dict[str, object]: components = build_slack_transcript_components(messages) text_body = str(components["text_content"] or "") if not text_body and placeholder_text: text_body = normalize_whitespace(placeholder_text) chat_metadata = { "author": None, "participants": components["participants"], "date_created": components["date_created"], "date_modified": components["date_modified"], "title": title, "message_count": components["message_count"], "timestamped_message_count": components["timestamped_message_count"], } payload = build_chat_extracted_payload( title=title, author=None, date_created=components["date_created"], text_body=text_body, html_body=None, attachments=[], preview_file_name=preview_file_name, chat_metadata=chat_metadata, chat_entries=list(components["chat_entries"]), ) if components.get("entity_hints"): payload["entity_hints"] = components["entity_hints"] return payload def build_slack_day_record( day_file: Path, *, rel_path: str, conversation_meta: dict[str, str], user_directory: dict[str, dict[str, str | None]], entity_identifier_scope: str | None = None, ) -> dict[str, object]: return { "day_file": day_file, "rel_path": rel_path, "folder_name": day_file.parent.name, "conversation_key": conversation_meta["conversation_key"], "conversation_type": conversation_meta["conversation_type"], "display_name": conversation_meta["display_name"], "title": slack_day_document_title(conversation_meta["display_name"], day_file), "messages": load_slack_day_messages( day_file, rel_path=rel_path, user_directory=user_directory, entity_identifier_scope=entity_identifier_scope, ), } def build_slack_thread_document_plans( conversation_key: str, day_records: list[dict[str, object]], ) -> tuple[list[dict[str, object]], set[str]]: day_records_by_rel = { str(record["rel_path"]): record for record in day_records } root_messages_by_ts: dict[str, dict[str, object]] = {} reply_messages_by_thread_ts: dict[str, list[dict[str, object]]] = defaultdict(list) candidate_thread_timestamps: set[str] = set() for day_record in day_records: for message in list(day_record.get("messages", [])): ts_raw = normalize_whitespace(str(message.get("ts") or "")) if not ts_raw: continue if bool(message.get("is_reply")): thread_ts = normalize_whitespace(str(message.get("thread_ts") or "")) if not thread_ts: continue reply_messages_by_thread_ts[thread_ts].append(message) candidate_thread_timestamps.add(thread_ts) continue root_messages_by_ts[ts_raw] = message if bool(message.get("is_thread_root")): candidate_thread_timestamps.add(ts_raw) thread_plans: list[dict[str, object]] = [] materialized_thread_roots: set[str] = set() for thread_ts in sorted(candidate_thread_timestamps, key=slack_timestamp_sort_key): root_message = root_messages_by_ts.get(thread_ts) if root_message is None: continue root_day_rel_path = str(root_message["day_rel_path"]) root_day_record = day_records_by_rel[root_day_rel_path] replies = sorted( reply_messages_by_thread_ts.get(thread_ts, []), key=lambda message: ( slack_timestamp_sort_key(message.get("ts")), int(message.get("ordinal") or 0), ), ) materialized_thread_roots.add(thread_ts) file_name = f"{thread_ts}.slackthread" source_parts = [ { "part_kind": "slack_thread_root_day", "rel_source_path": root_day_rel_path, "ordinal": 0, "label": root_day_record["title"], } ] seen_reply_days: set[str] = set() part_ordinal = 1 for reply in replies: reply_day_rel_path = normalize_whitespace(str(reply.get("day_rel_path") or "")) if not reply_day_rel_path or reply_day_rel_path == root_day_rel_path or reply_day_rel_path in seen_reply_days: continue seen_reply_days.add(reply_day_rel_path) reply_day_record = day_records_by_rel.get(reply_day_rel_path) source_parts.append( { "part_kind": "slack_thread_reply_day", "rel_source_path": reply_day_rel_path, "ordinal": part_ordinal, "label": reply_day_record["title"] if reply_day_record is not None else Path(reply_day_rel_path).name, } ) part_ordinal += 1 thread_plans.append( { "rel_path": slack_reply_thread_rel_path(conversation_key, thread_ts), "file_name": file_name, "preview_file_name": slack_preview_file_name(file_name), "title": slack_reply_thread_title( str(root_day_record["display_name"]), str(root_message.get("timestamp") or ""), ), "messages": [root_message, *replies], "parent_rel_path": root_day_rel_path, "source_rel_path": root_day_rel_path, "source_item_id": thread_ts, "source_folder_path": str(root_day_record["folder_name"]), "root_message_key": slack_root_message_key(conversation_key, thread_ts), "source_parts": source_parts, } ) return thread_plans, materialized_thread_roots def build_slack_day_document_plan( day_record: dict[str, object], *, materialized_thread_roots: set[str], ) -> dict[str, object]: visible_messages: list[dict[str, object]] = [] for message in list(day_record.get("messages", [])): if bool(message.get("is_reply")): thread_ts = normalize_whitespace(str(message.get("thread_ts") or "")) if thread_ts and thread_ts in materialized_thread_roots: continue visible_messages.append(message) file_name = str(day_record["day_file"].name) return { "rel_path": str(day_record["rel_path"]), "file_name": file_name, "preview_file_name": slack_preview_file_name(file_name), "title": str(day_record["title"]), "messages": visible_messages, "source_rel_path": str(day_record["rel_path"]), "source_item_id": None, "source_folder_path": str(day_record["folder_name"]), "root_message_key": None, "source_parts": [ { "part_kind": "slack_day_file", "rel_source_path": str(day_record["rel_path"]), "ordinal": 0, "label": day_record["title"], } ], } def plan_slack_export_conversations( root: Path, export_root: Path, *, conversation_directory: dict[str, dict[str, str]], user_directory: dict[str, dict[str, str | None]], day_files: list[Path], ) -> list[dict[str, object]]: rel_root = relative_document_path(root, export_root) conversation_day_records: dict[tuple[str, str, str], list[dict[str, object]]] = defaultdict(list) for day_file in day_files: rel_path = relative_document_path(root, day_file) folder_name = day_file.parent.name conversation_meta = slack_conversation_metadata_for_folder(folder_name, conversation_directory) conversation_identity = ( SLACK_EXPORT_SOURCE_KIND, rel_root, conversation_meta["conversation_key"], ) conversation_day_records[conversation_identity].append( build_slack_day_record( day_file, rel_path=rel_path, conversation_meta=conversation_meta, user_directory=user_directory, entity_identifier_scope=rel_root, ) ) conversation_plans: list[dict[str, object]] = [] for conversation_identity, records in sorted(conversation_day_records.items()): ordered_records = sorted(records, key=lambda record: (str(record["rel_path"]).lower(), str(record["rel_path"]))) if not ordered_records: continue thread_plans, materialized_thread_roots = build_slack_thread_document_plans( conversation_identity[2], ordered_records, ) day_documents: list[dict[str, object]] = [] rel_paths: list[str] = [] for day_record in ordered_records: day_plan = build_slack_day_document_plan( day_record, materialized_thread_roots=materialized_thread_roots, ) if not list(day_plan.get("messages") or []): continue day_documents.append( { "kind": "day", "plan": day_plan, "source_path": Path(day_record["day_file"]), } ) rel_paths.append(str(day_plan["rel_path"])) thread_documents: list[dict[str, object]] = [] for thread_plan in thread_plans: thread_documents.append( { "kind": "reply_thread", "plan": thread_plan, } ) rel_paths.append(str(thread_plan["rel_path"])) conversation_plans.append( { "conversation_identity": conversation_identity, "conversation_key": conversation_identity[2], "conversation_type": str(ordered_records[0]["conversation_type"]), "display_name": str(ordered_records[0]["display_name"]), "day_documents": day_documents, "thread_documents": thread_documents, "rel_paths": rel_paths, } ) return conversation_plans def prepare_slack_document_plan(document_plan: dict[str, object]) -> dict[str, object]: prepared_document = dict(document_plan) prepared_plan = dict(document_plan.get("plan") or {}) prepare_started = time.perf_counter() try: extracted_payload = build_slack_chat_payload( title=str(prepared_plan["title"]), preview_file_name=str(prepared_plan["preview_file_name"]), messages=list(prepared_plan.get("messages", [])), ) chunk_started = time.perf_counter() prepared_chunks = extracted_search_chunks(extracted_payload) prepared_document["plan"] = prepared_plan prepared_document["extracted_payload"] = extracted_payload prepared_document["prepared_chunks"] = prepared_chunks prepared_document["prepare_chunk_ms"] = (time.perf_counter() - chunk_started) * 1000.0 prepared_document["prepare_error"] = None except Exception as exc: prepared_document["plan"] = prepared_plan prepared_document["extracted_payload"] = None prepared_document["prepared_chunks"] = [] prepared_document["prepare_chunk_ms"] = 0.0 prepared_document["prepare_error"] = f"{type(exc).__name__}: {exc}" prepared_document["prepare_ms"] = (time.perf_counter() - prepare_started) * 1000.0 return prepared_document def prepare_slack_conversation_plan(conversation_plan: dict[str, object]) -> dict[str, object]: prepared_conversation = dict(conversation_plan) prepare_started = time.perf_counter() prepared_day_documents: list[dict[str, object]] = [] prepared_thread_documents: list[dict[str, object]] = [] prepare_error: str | None = None prepare_chunk_ms = 0.0 for document_plan in list(conversation_plan.get("day_documents") or []): prepared_document = prepare_slack_document_plan(document_plan) prepared_day_documents.append(prepared_document) prepare_chunk_ms += float(prepared_document.get("prepare_chunk_ms") or 0.0) if prepare_error is None and prepared_document.get("prepare_error"): prepare_error = f"{prepared_document['plan'].get('rel_path')}: {prepared_document['prepare_error']}" for document_plan in list(conversation_plan.get("thread_documents") or []): prepared_document = prepare_slack_document_plan(document_plan) prepared_thread_documents.append(prepared_document) prepare_chunk_ms += float(prepared_document.get("prepare_chunk_ms") or 0.0) if prepare_error is None and prepared_document.get("prepare_error"): prepare_error = f"{prepared_document['plan'].get('rel_path')}: {prepared_document['prepare_error']}" prepared_conversation["day_documents"] = prepared_day_documents prepared_conversation["thread_documents"] = prepared_thread_documents prepared_conversation["prepare_chunk_ms"] = prepare_chunk_ms prepared_conversation["prepare_error"] = prepare_error prepared_conversation["prepare_ms"] = (time.perf_counter() - prepare_started) * 1000.0 return prepared_conversation def iter_prepared_slack_conversation_plans( conversation_plans: list[dict[str, object]], staging_root: Path | None = None, ) -> Iterator[tuple[dict[str, object], float]]: effective_staging_root = staging_root if staging_root is not None and conversation_plans: effective_staging_root = ( Path(staging_root) / "slack" / sanitize_storage_filename(str(conversation_plans[0]["conversation_identity"][1])) ) yield from iter_staged_prepared_items( conversation_plans, prepare_item=prepare_slack_conversation_plan, config_benchmark_name="ingest_slack_prepare_config", queue_done_benchmark_name="ingest_slack_prepare_queue_done", spill_subdir_name="prepared-slack-conversations", staging_root=effective_staging_root, ) def remove_source_dataset_membership_for_document( connection: sqlite3.Connection, *, document_id: int, source_kind: str, source_locator: str, ) -> int: dataset_source_row = get_dataset_source_row( connection, source_kind=source_kind, source_locator=source_locator, ) if dataset_source_row is None: return 0 cursor = connection.execute( """ DELETE FROM dataset_documents WHERE document_id = ? AND dataset_source_id = ? """, (document_id, int(dataset_source_row["id"])), ) refresh_document_dataset_cache(connection, document_id) return int(cursor.rowcount or 0) def existing_rows_by_rel_path( connection: sqlite3.Connection, rel_paths: list[str], ) -> dict[str, sqlite3.Row]: if not rel_paths: return {} placeholders = ", ".join("?" for _ in rel_paths) rows = connection.execute( f""" SELECT * FROM documents WHERE rel_path IN ({placeholders}) ORDER BY id ASC """, rel_paths, ).fetchall() return {str(row["rel_path"]): row for row in rows} def existing_slack_document_row( connection: sqlite3.Connection, *, rel_path: str, source_rel_path: str, source_item_id: str | None, ) -> sqlite3.Row | None: row = connection.execute( """ SELECT * FROM documents WHERE rel_path = ? ORDER BY id ASC LIMIT 1 """, (rel_path,), ).fetchone() if row is not None: return row if source_item_id is None: return None return connection.execute( """ SELECT * FROM documents WHERE source_kind = ? AND source_rel_path = ? AND source_item_id = ? AND lifecycle_status != 'deleted' ORDER BY id ASC LIMIT 1 """, (SLACK_EXPORT_SOURCE_KIND, source_rel_path, source_item_id), ).fetchone() def commit_prepared_slack_day_document( connection: sqlite3.Connection, paths: dict[str, Path], prepared_document: dict[str, object], existing_row: sqlite3.Row | None, *, dataset_id: int, dataset_source_id: int, conversation_id: int, current_batch: int | None, ) -> dict[str, object]: plan = dict(prepared_document.get("plan") or {}) rel_path = str(plan["rel_path"]) effective_existing_row = existing_row or existing_slack_document_row( connection, rel_path=rel_path, source_rel_path=str(plan["source_rel_path"]), source_item_id=( str(plan["source_item_id"]) if plan.get("source_item_id") is not None else None ), ) extracted = apply_manual_locks(effective_existing_row, dict(prepared_document.get("extracted_payload") or {})) if effective_existing_row is None: if current_batch is None: current_batch = allocate_ingestion_batch_number(connection) control_number_batch = int(current_batch) control_number_family_sequence = reserve_control_number_family_sequence(connection, control_number_batch) control_number = format_control_number(control_number_batch, control_number_family_sequence) control_number_attachment_sequence = None else: control_number_batch = int(effective_existing_row["control_number_batch"]) control_number_family_sequence = int(effective_existing_row["control_number_family_sequence"]) control_number = str(effective_existing_row["control_number"]) control_number_attachment_sequence = effective_existing_row["control_number_attachment_sequence"] cleanup_document_artifacts(paths, connection, effective_existing_row) source_path_value = prepared_document.get("source_path") if source_path_value is None: raise RetrieverError(f"Slack day document {rel_path} is missing source_path.") document_id = upsert_document_row( connection, rel_path, Path(str(source_path_value)), effective_existing_row, extracted, file_name=str(plan["file_name"]), parent_document_id=None, control_number=control_number, dataset_id=dataset_id, conversation_id=conversation_id, control_number_batch=control_number_batch, control_number_family_sequence=control_number_family_sequence, control_number_attachment_sequence=control_number_attachment_sequence, source_kind=SLACK_EXPORT_SOURCE_KIND, source_rel_path=str(plan["source_rel_path"]), source_item_id=plan["source_item_id"], source_folder_path=str(plan["source_folder_path"]), ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=effective_existing_row, ) ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=dataset_source_id, ) remove_source_dataset_membership_for_document( connection, document_id=document_id, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), ) preview_rows = write_preview_artifacts(paths, rel_path, list(extracted.get("preview_artifacts", []))) replace_document_related_rows( connection, document_id, extracted | {"file_name": str(plan["file_name"])}, list(prepared_document.get("prepared_chunks") or []), preview_rows, ) replace_document_source_parts( connection, document_id, list(plan["source_parts"]), ) return { "document_id": int(document_id), "new": 1 if effective_existing_row is None else 0, "updated": 0 if effective_existing_row is None else 1, "current_batch": current_batch, "parent_state": { "document_id": int(document_id), "control_number_batch": int(control_number_batch), "control_number_family_sequence": int(control_number_family_sequence), }, } def commit_prepared_slack_thread_document( connection: sqlite3.Connection, paths: dict[str, Path], prepared_document: dict[str, object], existing_row: sqlite3.Row | None, *, dataset_id: int, dataset_source_id: int, conversation_id: int, parent_state_by_rel: dict[str, dict[str, int]] | None = None, ) -> dict[str, object]: plan = dict(prepared_document.get("plan") or {}) rel_path = str(plan["rel_path"]) parent_rel_path = str(plan["parent_rel_path"]) effective_existing_row = existing_row or existing_slack_document_row( connection, rel_path=rel_path, source_rel_path=str(plan["source_rel_path"]), source_item_id=( str(plan["source_item_id"]) if plan.get("source_item_id") is not None else None ), ) parent_state = dict((parent_state_by_rel or {}).get(parent_rel_path) or {}) if not parent_state: parent_row = connection.execute( """ SELECT id, control_number_batch, control_number_family_sequence FROM documents WHERE rel_path = ? AND lifecycle_status != 'deleted' ORDER BY id ASC LIMIT 1 """, (parent_rel_path,), ).fetchone() if parent_row is None: raise RetrieverError( f"Slack thread document {rel_path} is missing parent day document {parent_rel_path}." ) parent_state = { "document_id": int(parent_row["id"]), "control_number_batch": int(parent_row["control_number_batch"]), "control_number_family_sequence": int(parent_row["control_number_family_sequence"]), } parent_document_id = int(parent_state["document_id"]) extracted = apply_manual_locks(effective_existing_row, dict(prepared_document.get("extracted_payload") or {})) if effective_existing_row is None: control_number_batch = int(parent_state["control_number_batch"]) control_number_family_sequence = int(parent_state["control_number_family_sequence"]) control_number_attachment_sequence = next_attachment_sequence(connection, parent_document_id) control_number = format_control_number( control_number_batch, control_number_family_sequence, control_number_attachment_sequence, ) else: control_number_batch = int(effective_existing_row["control_number_batch"]) control_number_family_sequence = int(effective_existing_row["control_number_family_sequence"]) control_number_attachment_sequence = int(effective_existing_row["control_number_attachment_sequence"]) control_number = str(effective_existing_row["control_number"]) cleanup_document_artifacts(paths, connection, effective_existing_row) document_id = upsert_document_row( connection, rel_path, None, effective_existing_row, extracted, file_name=str(plan["file_name"]), parent_document_id=parent_document_id, child_document_kind=CHILD_DOCUMENT_KIND_REPLY_THREAD, control_number=control_number, dataset_id=dataset_id, conversation_id=conversation_id, control_number_batch=control_number_batch, control_number_family_sequence=control_number_family_sequence, control_number_attachment_sequence=control_number_attachment_sequence, root_message_key=str(plan["root_message_key"]), source_kind=SLACK_EXPORT_SOURCE_KIND, source_rel_path=str(plan["source_rel_path"]), source_item_id=str(plan["source_item_id"]), source_folder_path=str(plan["source_folder_path"]), ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=effective_existing_row, ) ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=dataset_source_id, ) remove_source_dataset_membership_for_document( connection, document_id=document_id, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), ) preview_rows = write_preview_artifacts(paths, rel_path, list(extracted.get("preview_artifacts", []))) replace_document_related_rows( connection, document_id, extracted | {"file_name": str(plan["file_name"])}, list(prepared_document.get("prepared_chunks") or []), preview_rows, ) replace_document_source_parts( connection, document_id, list(plan["source_parts"]), ) return { "document_id": int(document_id), "new": 1 if effective_existing_row is None else 0, "updated": 0 if effective_existing_row is None else 1, } def commit_prepared_slack_conversation( connection: sqlite3.Connection, paths: dict[str, Path], prepared_conversation: dict[str, object], existing_by_rel: dict[str, sqlite3.Row], *, dataset_id: int, dataset_source_id: int, current_batch: int | None, before_transaction_commit=None, ) -> dict[str, object]: prepare_error = prepared_conversation.get("prepare_error") if prepare_error: return { "status": "failed", "action": "failed", "current_batch": current_batch, "rel_paths": list(prepared_conversation.get("rel_paths") or []), "error": str(prepare_error), } rel_paths = list(prepared_conversation.get("rel_paths") or []) affected_document_ids: list[int] = [] parent_state_by_rel: dict[str, dict[str, int]] = {} connection.execute("BEGIN") try: conversation_id = upsert_conversation_row( connection, source_kind=SLACK_EXPORT_SOURCE_KIND, source_locator=str(prepared_conversation["conversation_identity"][1]), conversation_key=str(prepared_conversation["conversation_key"]), conversation_type=str(prepared_conversation["conversation_type"]), display_name=str(prepared_conversation["display_name"]), ) new_count = 0 updated_count = 0 for prepared_document in list(prepared_conversation.get("day_documents") or []): plan = dict(prepared_document.get("plan") or {}) rel_path = str(plan["rel_path"]) commit_payload = commit_prepared_slack_day_document( connection, paths, prepared_document, existing_by_rel.get(rel_path), dataset_id=dataset_id, dataset_source_id=dataset_source_id, conversation_id=conversation_id, current_batch=current_batch, ) current_batch = commit_payload["current_batch"] affected_document_ids.append(int(commit_payload["document_id"])) parent_state_by_rel[rel_path] = dict(commit_payload["parent_state"]) new_count += int(commit_payload["new"]) updated_count += int(commit_payload["updated"]) for prepared_document in list(prepared_conversation.get("thread_documents") or []): plan = dict(prepared_document.get("plan") or {}) rel_path = str(plan["rel_path"]) commit_payload = commit_prepared_slack_thread_document( connection, paths, prepared_document, existing_by_rel.get(rel_path), dataset_id=dataset_id, dataset_source_id=dataset_source_id, conversation_id=conversation_id, parent_state_by_rel=parent_state_by_rel, ) affected_document_ids.append(int(commit_payload["document_id"])) new_count += int(commit_payload["new"]) updated_count += int(commit_payload["updated"]) result = { "status": "ok", "action": "committed", "current_batch": current_batch, "new": new_count, "updated": updated_count, "affected_document_ids": affected_document_ids, "rel_paths": rel_paths, "source_locator": str(prepared_conversation["conversation_identity"][1]), "conversation_key": str(prepared_conversation["conversation_key"]), } if before_transaction_commit is not None: before_transaction_commit(connection, result) connection.commit() return result except Exception as exc: connection.rollback() return { "status": "failed", "action": "failed", "current_batch": current_batch, "rel_paths": rel_paths, "error": f"{type(exc).__name__}: {exc}", } def commit_prepared_slack_document_in_transaction( connection: sqlite3.Connection, paths: dict[str, Path], prepared_document: dict[str, object], *, dataset_id: int, dataset_source_id: int, current_batch: int | None, ) -> dict[str, object]: prepare_error = prepared_document.get("prepare_error") plan = dict(prepared_document.get("plan") or {}) rel_path = str(plan.get("rel_path") or "") if prepare_error: return { "status": "failed", "action": "failed", "current_batch": current_batch, "rel_paths": [rel_path] if rel_path else [], "error": str(prepare_error), } source_locator = str(prepared_document.get("source_locator") or "") conversation_key = str(prepared_document.get("conversation_key") or "") conversation_type = str(prepared_document.get("conversation_type") or "") display_name = str(prepared_document.get("display_name") or "") document_kind = str(prepared_document.get("document_kind") or prepared_document.get("kind") or "") conversation_id = upsert_conversation_row( connection, source_kind=SLACK_EXPORT_SOURCE_KIND, source_locator=source_locator, conversation_key=conversation_key, conversation_type=conversation_type, display_name=display_name, ) existing_row = connection.execute( """ SELECT * FROM documents WHERE rel_path = ? ORDER BY id ASC LIMIT 1 """, (rel_path,), ).fetchone() if document_kind == "day": commit_payload = commit_prepared_slack_day_document( connection, paths, prepared_document, existing_row, dataset_id=dataset_id, dataset_source_id=dataset_source_id, conversation_id=conversation_id, current_batch=current_batch, ) current_batch = commit_payload["current_batch"] elif document_kind == "reply_thread": commit_payload = commit_prepared_slack_thread_document( connection, paths, prepared_document, existing_row, dataset_id=dataset_id, dataset_source_id=dataset_source_id, conversation_id=conversation_id, ) else: raise RetrieverError(f"Unsupported Slack document kind: {document_kind or ''}") return { "status": "ok", "action": "committed", "current_batch": current_batch, "new": int(commit_payload["new"]), "updated": int(commit_payload["updated"]), "affected_document_ids": [int(commit_payload["document_id"])], "rel_paths": [rel_path] if rel_path else [], "source_locator": source_locator, "conversation_key": conversation_key, "document_kind": document_kind, "conversation_lead_document": bool(prepared_document.get("conversation_lead_document")), } def commit_prepared_slack_document( connection: sqlite3.Connection, paths: dict[str, Path], prepared_document: dict[str, object], *, dataset_id: int, dataset_source_id: int, current_batch: int | None, before_transaction_commit=None, ) -> dict[str, object]: plan = dict(prepared_document.get("plan") or {}) rel_path = str(plan.get("rel_path") or "") current_batch_value = current_batch connection.execute("BEGIN") try: result = commit_prepared_slack_document_in_transaction( connection, paths, prepared_document, dataset_id=dataset_id, dataset_source_id=dataset_source_id, current_batch=current_batch_value, ) if result.get("current_batch") is not None: current_batch_value = result.get("current_batch") if before_transaction_commit is not None: before_transaction_commit(connection, result) connection.commit() return result except Exception as exc: connection.rollback() return { "status": "failed", "action": "failed", "current_batch": current_batch_value, "rel_paths": [rel_path] if rel_path else [], "error": f"{type(exc).__name__}: {exc}", } def mark_missing_slack_export_documents( connection: sqlite3.Connection, *, dataset_id: int, seen_rel_paths: set[str], ) -> int: rows = connection.execute( """ SELECT d.id, d.rel_path, d.lifecycle_status FROM documents d JOIN dataset_documents dd ON dd.document_id = d.id WHERE d.source_kind = ? AND dd.dataset_id = ? AND d.lifecycle_status != 'deleted' ORDER BY d.id ASC """, (SLACK_EXPORT_SOURCE_KIND, dataset_id), ).fetchall() missing_ids = [ int(row["id"]) for row in rows if normalize_whitespace(str(row["rel_path"] or "")) not in seen_rel_paths and row["lifecycle_status"] != "missing" ] if not missing_ids: return 0 now = utc_now() placeholders = ", ".join("?" for _ in missing_ids) connection.execute( f""" UPDATE documents SET lifecycle_status = 'missing', updated_at = ? WHERE lifecycle_status != 'deleted' AND id IN ({placeholders}) """, [now, *missing_ids], ) return len(missing_ids) def ingest_slack_export_root( connection: sqlite3.Connection, paths: dict[str, Path], export_root: Path, *, ingestion_batch_number: int | None = None, staging_root: Path | None = None, ) -> dict[str, object]: root = paths["root"] rel_root = relative_document_path(root, export_root) dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=SLACK_EXPORT_SOURCE_KIND, source_locator=rel_root, dataset_name=slack_export_dataset_name(rel_root), ) connection.commit() conversation_directory = load_slack_export_conversation_directory(export_root) user_directory = load_slack_user_directory(export_root) day_files = iter_slack_export_day_files(export_root) stats = { "new": 0, "updated": 0, "missing": 0, "failed": 0, "scanned_day_files": len(day_files), "conversations": 0, } failures: list[dict[str, str]] = [] current_batch = ingestion_batch_number conversation_plans = plan_slack_export_conversations( root, export_root, conversation_directory=conversation_directory, user_directory=user_directory, day_files=day_files, ) stats["conversations"] = len(conversation_plans) existing_by_rel = existing_rows_by_rel_path( connection, [rel_path for conversation_plan in conversation_plans for rel_path in list(conversation_plan.get("rel_paths") or [])], ) seen_rel_paths: set[str] = { rel_path for conversation_plan in conversation_plans for rel_path in list(conversation_plan.get("rel_paths") or []) } prepare_ms = 0.0 prepare_chunk_ms = 0.0 prepare_wait_ms = 0.0 commit_ms = 0.0 conversation_loop_started = time.perf_counter() for prepared_conversation, wait_ms in iter_prepared_slack_conversation_plans( conversation_plans, staging_root=staging_root, ): prepare_wait_ms += wait_ms prepare_ms += float(prepared_conversation.get("prepare_ms") or 0.0) prepare_chunk_ms += float(prepared_conversation.get("prepare_chunk_ms") or 0.0) commit_started = time.perf_counter() commit_result = commit_prepared_slack_conversation( connection, paths, prepared_conversation, existing_by_rel, dataset_id=dataset_id, dataset_source_id=dataset_source_id, current_batch=current_batch, ) commit_ms += (time.perf_counter() - commit_started) * 1000.0 current_batch = commit_result["current_batch"] if commit_result["status"] == "failed": stats["failed"] += 1 failures.append( { "rel_path": ", ".join(commit_result.get("rel_paths") or list(prepared_conversation.get("rel_paths") or [])), "error": str(commit_result.get("error") or "Unknown slack export ingest failure."), } ) continue stats["new"] += int(commit_result["new"]) stats["updated"] += int(commit_result["updated"]) benchmark_mark( "ingest_slack_conversations_done", conversation_count=len(conversation_plans), conversation_loop_ms=round((time.perf_counter() - conversation_loop_started) * 1000.0, 3), prepare_ms=round(prepare_ms, 3), prepare_chunk_ms=round(prepare_chunk_ms, 3), prepare_wait_ms=round(prepare_wait_ms, 3), commit_ms=round(commit_ms, 3), new=stats["new"], updated=stats["updated"], failed=stats["failed"], ) missing_started = time.perf_counter() connection.execute("BEGIN") try: stats["missing"] = mark_missing_slack_export_documents( connection, dataset_id=dataset_id, seen_rel_paths=seen_rel_paths, ) connection.commit() except Exception: connection.rollback() raise benchmark_mark( "ingest_slack_missing_done", missing_ms=round((time.perf_counter() - missing_started) * 1000.0, 3), missing=stats["missing"], ) return { **stats, "dataset_id": dataset_id, "dataset_source_id": dataset_source_id, "source_locator": rel_root, "ingestion_batch_number": current_batch, "failures": failures, } GMAIL_EXPORT_ARCHIVE_BROWSER_FILE = "archive_browser.html" GMAIL_EXPORT_DRIVE_FOLDER_PATTERN = re.compile(r"_Drive_Link_Export_\d+$", re.IGNORECASE) GMAIL_EXPORT_FILE_DRIVE_ID_PATTERN = re.compile(r"^(?P.+)_(?P[A-Za-z0-9_-]{10,})$") GMAIL_METADATA_REQUIRED_HEADERS = { "Rfc822MessageId", "GmailMessageId", "Labels", "Account", "Subject", } GMAIL_DRIVE_LINKS_REQUIRED_HEADERS = { "Rfc822MessageId", "DriveUrl", "DriveItemId", } GMAIL_RESULT_COUNTS_REQUIRED_HEADERS = { "Email", "AccountStatus", "SuccessCount", "MessageErrorCount", "ChatErrorCount", } PST_EXPORT_STANDARD_RESULTS_REQUIRED_HEADERS = { "Document ID", "Item Identity", "Target Path", } PST_EXPORT_ARCHIVE_RESULTS_REQUIRED_HEADERS = { "Document ID", "Export Item Id", "Export Item Path", } PST_EXPORT_RESULTS_FILE_PATTERN = re.compile(r"^results\.csv$", re.IGNORECASE) PST_EXPORT_ARCHIVE_RESULTS_FILE_PATTERN = re.compile(r"^export_results.*\.csv$", re.IGNORECASE) PST_EXPORT_SUMMARY_FILE_PATTERN = re.compile(r"^export summary .+\.csv$", re.IGNORECASE) PST_EXPORT_MANIFEST_FILE_PATTERN = re.compile(r"^manifest\.xml$", re.IGNORECASE) PST_EXPORT_TRACE_LOG_FILE_PATTERN = re.compile(r"^trace\.log$", re.IGNORECASE) def gmail_result_counts_csv_valid(path: Path) -> bool: headers, _ = load_normalized_csv_rows(path) normalized_headers = { normalize_inline_whitespace(str(header or "")).casefold() for header in headers if header } required_headers = {header.casefold() for header in GMAIL_RESULT_COUNTS_REQUIRED_HEADERS} return required_headers.issubset(normalized_headers) def gmail_export_family_keys(value: Path | str) -> set[str]: raw_name = value.stem if isinstance(value, Path) else str(value or "") pending = [normalize_whitespace(raw_name)] seen_candidates: set[str] = set() normalized_keys: set[str] = set() while pending: candidate = normalize_whitespace(pending.pop()) candidate_key = candidate.casefold() if not candidate or candidate_key in seen_candidates: continue seen_candidates.add(candidate_key) normalized_key = re.sub(r"[^a-z0-9]+", "", candidate_key) if normalized_key: normalized_keys.add(normalized_key) if "--" in candidate: pending.append(candidate.split("--", 1)[0]) if candidate_key.endswith("-errors"): pending.append(candidate[: -len("-errors")]) if candidate_key.endswith("-result-counts"): pending.append(candidate[: -len("-result-counts")]) trimmed_numeric_suffix = re.sub(r"[-_ ]+\d+$", "", candidate) if trimmed_numeric_suffix != candidate: pending.append(trimmed_numeric_suffix) return normalized_keys def load_normalized_csv_rows(path: Path) -> tuple[list[str], list[dict[str, str]]]: try: decoded, _, _ = decode_bytes(path.read_bytes()) except OSError: return [], [] try: reader = csv.DictReader(io.StringIO(decoded)) except csv.Error: return [], [] headers = [normalize_inline_whitespace(str(header or "")) for header in (reader.fieldnames or []) if header] rows: list[dict[str, str]] = [] for raw_row in reader: if not isinstance(raw_row, dict): continue normalized_row: dict[str, str] = {} for key, value in raw_row.items(): normalized_key = normalize_inline_whitespace(str(key or "")) if not normalized_key: continue normalized_row[normalized_key] = normalize_whitespace(str(value or "")) if value is not None else "" if normalized_row: rows.append(normalized_row) return headers, rows def csv_has_required_headers(path: Path, required_headers: set[str]) -> bool: headers, _ = load_normalized_csv_rows(path) return required_headers.issubset(set(headers)) def load_gmail_csv_rows(path: Path) -> tuple[list[str], list[dict[str, str]]]: return load_normalized_csv_rows(path) def gmail_csv_has_required_headers(path: Path, required_headers: set[str]) -> bool: return csv_has_required_headers(path, required_headers) def gmail_metadata_csv_valid(path: Path) -> bool: return path.is_file() and gmail_csv_has_required_headers(path, GMAIL_METADATA_REQUIRED_HEADERS) def gmail_drive_links_csv_valid(path: Path) -> bool: return path.is_file() and gmail_csv_has_required_headers(path, GMAIL_DRIVE_LINKS_REQUIRED_HEADERS) def gmail_csv_list_values(value: object) -> list[str]: values: list[str] = [] seen: set[str] = set() for raw_part in str(value or "").split(","): normalized = normalize_whitespace(raw_part) if not normalized or normalized in seen: continue values.append(normalized) seen.add(normalized) return values def gmail_normalized_message_lookup_key(value: object) -> str | None: return normalize_email_message_id(value) def gmail_normalized_drive_item_id(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")) return normalized or None def gmail_drive_item_id_from_export_file_name(path: Path) -> str | None: stem = normalize_whitespace(path.stem) if "_" not in stem: return None match = GMAIL_EXPORT_FILE_DRIVE_ID_PATTERN.fullmatch(stem) if match is None: return None return gmail_normalized_drive_item_id(match.group("drive_id")) def gmail_drive_export_files(export_dirs: list[Path]) -> list[Path]: files: list[Path] = [] for export_dir in export_dirs: try: iterator = export_dir.rglob("*") except OSError: continue for path in iterator: if path.is_dir() or ".retriever" in path.parts: continue file_type = normalize_extension(path) if not file_type or file_type not in SUPPORTED_FILE_TYPES: continue files.append(path) return sorted(files) def parse_gmail_metadata_csv(paths: list[Path]) -> dict[str, dict[str, object]]: metadata_by_message_id: dict[str, dict[str, object]] = {} for path in paths: _, rows = load_gmail_csv_rows(path) for row in rows: message_id = gmail_normalized_message_lookup_key(row.get("Rfc822MessageId")) if not message_id: continue recipients = ", ".join( part for part in [row.get("To"), row.get("CC"), row.get("BCC")] if normalize_whitespace(str(part or "")) ) metadata_by_message_id[message_id] = { "account": normalize_whitespace(str(row.get("Account") or "")) or None, "bcc": normalize_whitespace(str(row.get("BCC") or "")) or None, "cc": normalize_whitespace(str(row.get("CC") or "")) or None, "date_received": normalize_date_field_value(row.get("DateReceived")), "date_sent": normalize_date_field_value(row.get("DateSent")), "file_name": normalize_whitespace(str(row.get("FileName") or "")) or None, "from": normalize_whitespace(str(row.get("From") or "")) or None, "gmail_message_id": gmail_normalized_drive_item_id(row.get("GmailMessageId")), "labels": gmail_csv_list_values(row.get("Labels")), "subject": normalize_generated_document_title(row.get("Subject")), "threaded_message_count": normalize_whitespace(str(row.get("ThreadedMessageCount") or "")) or None, "to": normalize_whitespace(str(row.get("To") or "")) or None, "recipients": recipients or None, } return metadata_by_message_id def parse_gmail_drive_links_csv(paths: list[Path]) -> dict[str, list[dict[str, str]]]: links_by_message_id: dict[str, list[dict[str, str]]] = defaultdict(list) seen: dict[str, set[tuple[str, str]]] = defaultdict(set) for path in paths: _, rows = load_gmail_csv_rows(path) for row in rows: message_id = gmail_normalized_message_lookup_key(row.get("Rfc822MessageId")) drive_item_id = gmail_normalized_drive_item_id(row.get("DriveItemId")) if not message_id or not drive_item_id: continue drive_url = html.unescape(normalize_whitespace(str(row.get("DriveUrl") or ""))) or "" marker = (drive_item_id, drive_url) if marker in seen[message_id]: continue links_by_message_id[message_id].append( { "drive_item_id": drive_item_id, "drive_url": drive_url, } ) seen[message_id].add(marker) return links_by_message_id def parse_gmail_drive_export_metadata( paths: list[Path], *, file_paths_by_name: dict[str, Path], ) -> dict[str, dict[str, object]]: metadata_by_drive_item_id: dict[str, dict[str, object]] = {} for path in paths: try: root = parse_xml_document(path.read_bytes()) except (OSError, ET.ParseError): continue for document in root.findall(".//Document"): drive_item_id = gmail_normalized_drive_item_id(document.attrib.get("DocID")) if not drive_item_id: continue tag_values: dict[str, str] = {} for tag in document.findall("./Tags/Tag"): tag_name = normalize_inline_whitespace(str(tag.attrib.get("TagName") or "")) if not tag_name: continue tag_values[tag_name] = normalize_whitespace(str(tag.attrib.get("TagValue") or "")) external_file = document.find(".//ExternalFile") file_name = normalize_whitespace(str(external_file.attrib.get("FileName") or "")) if external_file is not None else "" file_path = file_paths_by_name.get(file_name) metadata_by_drive_item_id[drive_item_id] = { "author": normalize_whitespace(str(tag_values.get("#Author") or "")) or None, "collaborators": gmail_csv_list_values(tag_values.get("Collaborators")), "date_created": normalize_date_field_value(tag_values.get("#DateCreated")), "date_modified": normalize_date_field_value(tag_values.get("#DateModified")), "document_type": normalize_whitespace(str(tag_values.get("DocumentType") or "")) or None, "drive_item_id": drive_item_id, "file_name": file_name or (file_path.name if file_path is not None else None), "file_path": file_path, "others": gmail_csv_list_values(tag_values.get("Others")), "source_hash": normalize_whitespace(str(tag_values.get("SourceHash") or "")) or None, "title": normalize_generated_document_title(tag_values.get("#Title")) or None, "viewers": gmail_csv_list_values(tag_values.get("Viewers")), } return metadata_by_drive_item_id def parse_gmail_drive_export_errors(paths: list[Path]) -> dict[str, str]: errors_by_drive_item_id: dict[str, str] = {} for path in paths: _, rows = load_gmail_csv_rows(path) for row in rows: drive_item_id = ( gmail_normalized_drive_item_id(row.get("Drive Document ID")) or gmail_normalized_drive_item_id(row.get("Document ID")) ) if not drive_item_id: continue description = normalize_whitespace(str(row.get("Error Description") or "")) if description: errors_by_drive_item_id[drive_item_id] = description return errors_by_drive_item_id def gmail_drive_document_participants(record: dict[str, object]) -> list[str]: values: list[str] = [] for value in [ record.get("author"), *list(record.get("collaborators") or []), *list(record.get("viewers") or []), *list(record.get("others") or []), ]: normalized = normalize_whitespace(str(value or "")) if normalized: values.append(normalized) return sorted_unique_display_names(values) def gmail_drive_document_title(record: dict[str, object]) -> str | None: for value in ( record.get("title"), record.get("file_name"), ): normalized = normalize_generated_document_title(value) if normalized: return normalized return None def gmail_drive_document_link_summary(record: dict[str, object]) -> dict[str, object]: return { "author": normalize_whitespace(str(record.get("author") or "")) or None, "drive_item_id": gmail_normalized_drive_item_id(record.get("drive_item_id")), "error": normalize_whitespace(str(record.get("error") or "")) or None, "file_name": normalize_whitespace(str(record.get("file_name") or "")) or None, "title": gmail_drive_document_title(record), } def gmail_drive_attachment_payload(record: dict[str, object]) -> dict[str, object] | None: file_path = record.get("file_path") if not isinstance(file_path, Path) or not file_path.exists() or file_path.is_dir(): return None try: payload = file_path.read_bytes() except OSError: return None file_name = normalize_whitespace(str(record.get("file_name") or "")) or file_path.name return { "file_name": file_name, "payload": payload, "file_hash": sha256_bytes(payload), "gmail_drive_record": dict(record), } def gmail_drive_record_preference_key(record: dict[str, object]) -> tuple[int, int, int, int, int, int]: drive_item_id = normalize_whitespace(str(record.get("drive_item_id") or "")) return ( 1 if list(record.get("linked_message_ids") or []) else 0, 1 if gmail_drive_document_title(record) else 0, 1 if normalize_whitespace(str(record.get("author") or "")) else 0, 1 if normalize_whitespace(str(record.get("date_created") or "")) else 0, 1 if normalize_whitespace(str(record.get("date_modified") or "")) else 0, len(drive_item_id), ) def append_extracted_search_context( extracted: dict[str, object], extra_sections: list[str], ) -> dict[str, object]: normalized_sections = [ normalize_whitespace(section) for section in extra_sections if normalize_whitespace(section) ] if not normalized_sections: return dict(extracted) merged = dict(extracted) chunks = list(extracted_search_chunks(merged)) next_chunk_index = max((int(chunk["chunk_index"]) for chunk in chunks), default=-1) + 1 for section in normalized_sections: chunks.append( { "chunk_index": next_chunk_index, "char_start": 0, "char_end": len(section), "token_estimate": token_estimate(section), "text_content": section, } ) next_chunk_index += 1 merged["chunks"] = chunks if merged.get("text_status") == "empty": merged["text_status"] = "ok" return merged def gmail_append_search_context( extracted: dict[str, object], extra_sections: list[str], ) -> dict[str, object]: return append_extracted_search_context(extracted, extra_sections) def gmail_email_metadata_search_text( message_metadata: dict[str, object] | None, linked_drive_records: list[dict[str, object]], ) -> str | None: lines: list[str] = [] if message_metadata: lines.append("Gmail export metadata") account = normalize_whitespace(str(message_metadata.get("account") or "")) if account: lines.append(f"Account: {account}") gmail_message_id = normalize_whitespace(str(message_metadata.get("gmail_message_id") or "")) if gmail_message_id: lines.append(f"Gmail message ID: {gmail_message_id}") labels = [label for label in list(message_metadata.get("labels") or []) if normalize_whitespace(str(label or ""))] if labels: lines.append(f"Labels: {', '.join(labels)}") date_received = normalize_whitespace(str(message_metadata.get("date_received") or "")) if date_received: lines.append(f"Date received: {date_received}") threaded_count = normalize_whitespace(str(message_metadata.get("threaded_message_count") or "")) if threaded_count: lines.append(f"Threaded message count: {threaded_count}") if linked_drive_records: lines.append("Linked Google Drive items") for record in linked_drive_records[:20]: title = normalize_whitespace(str(record.get("title") or record.get("file_name") or record.get("drive_item_id") or "")) if not title: continue details = [ f"Drive ID: {record['drive_item_id']}" for _ in [record.get("drive_item_id")] if normalize_whitespace(str(record.get("drive_item_id") or "")) ] author = normalize_whitespace(str(record.get("author") or "")) if author: details.append(f"Author: {author}") error_text = normalize_whitespace(str(record.get("error") or "")) if error_text: details.append(f"Export error: {error_text}") lines.append(f"- {title}" + (f" ({'; '.join(details)})" if details else "")) normalized = normalize_whitespace("\n".join(lines)) return normalized or None def gmail_drive_document_search_text(record: dict[str, object]) -> str | None: lines: list[str] = ["Gmail Drive export metadata"] drive_item_id = normalize_whitespace(str(record.get("drive_item_id") or "")) if drive_item_id: lines.append(f"Drive item ID: {drive_item_id}") title = gmail_drive_document_title(record) if title: lines.append(f"Title: {title}") author = normalize_whitespace(str(record.get("author") or "")) if author: lines.append(f"Author: {author}") collaborators = [value for value in list(record.get("collaborators") or []) if normalize_whitespace(str(value or ""))] if collaborators: lines.append(f"Collaborators: {', '.join(collaborators)}") viewers = [value for value in list(record.get("viewers") or []) if normalize_whitespace(str(value or ""))] if viewers: lines.append(f"Viewers: {', '.join(viewers)}") others = [value for value in list(record.get("others") or []) if normalize_whitespace(str(value or ""))] if others: lines.append(f"Others: {', '.join(others)}") linked_subjects = [ normalize_generated_document_title(subject) for subject in list(record.get("linked_subjects") or []) if normalize_generated_document_title(subject) ] if linked_subjects: lines.append("Linked from Gmail messages") for subject in linked_subjects[:20]: lines.append(f"- {subject}") error_text = normalize_whitespace(str(record.get("error") or "")) if error_text: lines.append(f"Export error: {error_text}") normalized = normalize_whitespace("\n".join(lines)) return normalized or None def apply_gmail_email_export_metadata( extracted: dict[str, object], *, message_metadata: dict[str, object] | None, linked_drive_records: list[dict[str, object]], ) -> dict[str, object]: enriched = dict(extracted) if message_metadata: if not normalize_whitespace(str(enriched.get("date_created") or "")) and message_metadata.get("date_sent"): enriched["date_created"] = message_metadata.get("date_sent") subject = normalize_generated_document_title(message_metadata.get("subject")) if subject and not normalize_whitespace(str(enriched.get("title") or "")): enriched["title"] = subject if subject and not normalize_whitespace(str(enriched.get("subject") or "")): enriched["subject"] = subject author = normalize_whitespace(str(message_metadata.get("from") or "")) if author and not normalize_whitespace(str(enriched.get("author") or "")): enriched["author"] = author recipients = normalize_whitespace(str(message_metadata.get("recipients") or "")) if recipients and not normalize_whitespace(str(enriched.get("recipients") or "")): enriched["recipients"] = recipients search_text = gmail_email_metadata_search_text(message_metadata, linked_drive_records) return gmail_append_search_context(enriched, [search_text] if search_text else []) def apply_gmail_drive_export_metadata( extracted: dict[str, object], *, drive_record: dict[str, object], ) -> dict[str, object]: enriched = dict(extracted) title = gmail_drive_document_title(drive_record) if title: enriched["title"] = title author = normalize_whitespace(str(drive_record.get("author") or "")) if author: enriched["author"] = author date_created = normalize_whitespace(str(drive_record.get("date_created") or "")) if date_created: enriched["date_created"] = date_created date_modified = normalize_whitespace(str(drive_record.get("date_modified") or "")) if date_modified: enriched["date_modified"] = date_modified participants = gmail_drive_document_participants(drive_record) if participants: enriched["participants"] = "; ".join(participants) search_text = gmail_drive_document_search_text(drive_record) return gmail_append_search_context(enriched, [search_text] if search_text else []) def gmail_enriched_message_file_hash( base_hash: object, *, message_metadata: dict[str, object] | None, linked_drive_records: list[dict[str, object]], linked_drive_attachment_records: list[dict[str, object]] | None = None, ) -> str: return sha256_json_value( { "base_hash": normalize_whitespace(str(base_hash or "")) or None, "message_metadata": message_metadata or {}, "linked_drive_records": linked_drive_records, "linked_drive_attachments": [ { "drive_item_id": gmail_normalized_drive_item_id(record.get("drive_item_id")), "file_name": normalize_whitespace(str(record.get("file_name") or "")) or None, "file_hash": ( gmail_drive_document_file_hash(file_path, record) if isinstance(file_path := record.get("file_path"), Path) and file_path.exists() else None ), "title": gmail_drive_document_title(record), } for record in list(linked_drive_attachment_records or []) ], } ) def gmail_drive_document_file_hash(path: Path, drive_record: dict[str, object]) -> str: return sha256_json_value( { "file_hash": sha256_file(path), "drive_record": { "author": drive_record.get("author"), "collaborators": list(drive_record.get("collaborators") or []), "date_created": drive_record.get("date_created"), "date_modified": drive_record.get("date_modified"), "drive_item_id": drive_record.get("drive_item_id"), "error": drive_record.get("error"), "linked_message_ids": list(drive_record.get("linked_message_ids") or []), "linked_subjects": list(drive_record.get("linked_subjects") or []), "others": list(drive_record.get("others") or []), "title": drive_record.get("title"), "viewers": list(drive_record.get("viewers") or []), }, } ) def pst_export_results_csv_valid(path: Path) -> bool: if not path.is_file(): return False if PST_EXPORT_RESULTS_FILE_PATTERN.match(path.name): return csv_has_required_headers(path, PST_EXPORT_STANDARD_RESULTS_REQUIRED_HEADERS) if PST_EXPORT_ARCHIVE_RESULTS_FILE_PATTERN.match(path.name): return csv_has_required_headers(path, PST_EXPORT_ARCHIVE_RESULTS_REQUIRED_HEADERS) return False def pst_export_summary_csv_valid(path: Path) -> bool: return path.is_file() and PST_EXPORT_SUMMARY_FILE_PATTERN.match(path.name) is not None def pst_export_manifest_xml_valid(path: Path) -> bool: return path.is_file() and PST_EXPORT_MANIFEST_FILE_PATTERN.match(path.name) is not None def pst_export_trace_log_valid(path: Path) -> bool: return path.is_file() and PST_EXPORT_TRACE_LOG_FILE_PATTERN.match(path.name) is not None def pst_export_normalized_text(value: object) -> str | None: normalized = normalize_whitespace(html.unescape(str(value or ""))) return normalized or None PST_EXPORT_SMTP_ADDRESS_PATTERN = re.compile( r"[A-Za-z0-9.!#$%&'*+/=?^_`{|}~-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}" ) PST_EXPORT_EXCHANGE_DN_NOISE_PATTERN = re.compile( r"(?i)(?:^|\s)(?:/O=|/OU=|/CN=|=GUID-[0-9A-F-]{8,})" ) def pst_export_clean_custodian_candidate(value: object) -> str | None: candidate = pst_export_normalized_text(value) if not candidate: return None direct_email = normalize_entity_email(candidate) if direct_email: return direct_email emails: list[str] = [] for match in PST_EXPORT_SMTP_ADDRESS_PATTERN.finditer(candidate): email = normalize_entity_email(match.group(0)) if email and email not in emails: emails.append(email) if len(emails) == 1 and PST_EXPORT_EXCHANGE_DN_NOISE_PATTERN.search(candidate): return emails[0] return candidate def pst_export_match_text_key(value: object) -> str | None: normalized = normalize_generated_document_title(value) or pst_export_normalized_text(value) return normalized.casefold() if normalized else None def pst_export_row_value(row: dict[str, object], *keys: str) -> str | None: for key in keys: if key not in row: continue normalized = pst_export_normalized_text(row.get(key)) if normalized: return normalized return None def pst_export_path_parts(value: object) -> list[str]: normalized = pst_export_normalized_text(value) if not normalized: return [] return [part.strip() for part in re.split(r"[\\/]+", normalized) if part.strip()] def pst_export_folder_match_part(value: object) -> str | None: normalized = pst_export_normalized_text(value) if not normalized: return None normalized = normalize_whitespace(re.sub(r"[-_]+", " ", normalized.casefold())) return normalized or None def pst_export_folder_match_parts(value: object) -> tuple[str, ...]: parts = pst_export_path_parts(value) if not parts: return () pst_index = next((index for index, part in enumerate(parts) if part.lower().endswith(".pst")), None) if pst_index is not None: parts = parts[pst_index + 1 :] if parts: parts = parts[:-1] normalized_parts = [ normalized_part for normalized_part in (pst_export_folder_match_part(part) for part in parts) if normalized_part ] return tuple(normalized_parts) def pst_export_folder_parts_match( message_folder_parts: tuple[str, ...], candidate_folder_parts: tuple[str, ...], ) -> bool: if not message_folder_parts or not candidate_folder_parts: return False if len(message_folder_parts) >= len(candidate_folder_parts): return tuple(message_folder_parts[-len(candidate_folder_parts) :]) == tuple(candidate_folder_parts) return tuple(candidate_folder_parts[-len(message_folder_parts) :]) == tuple(message_folder_parts) def pst_export_resolve_pst_path( candidate_root: Path, raw_path: object, ) -> tuple[Path | None, bool]: parts = pst_export_path_parts(raw_path) pst_index = next((index for index, part in enumerate(parts) if part.lower().endswith(".pst")), None) if pst_index is None: return None, False has_message_component = pst_index < len(parts) - 1 pst_name = parts[pst_index] candidate_part_sets: list[list[str]] = [] if parts and parts[0].lower() == "exchange": candidate_part_sets.append(parts[: pst_index + 1]) candidate_part_sets.extend((["Exchange", pst_name], [pst_name])) seen: set[tuple[str, ...]] = set() for rel_parts in candidate_part_sets: marker = tuple(part.lower() for part in rel_parts) if marker in seen: continue seen.add(marker) resolved = resolve_case_insensitive_relative_path(candidate_root, rel_parts) if resolved is None or not resolved.exists() or resolved.is_dir(): continue if normalize_extension(resolved) != PST_SOURCE_KIND: continue return resolved.resolve(), has_message_component return None, has_message_component def pst_export_combined_recipients(row: dict[str, object]) -> str | None: recipients: list[str] = [] for keys in ( ("To - Expanded", "To – Expanded", "Recipients in To line"), ("CC - Expanded", "CC – Expanded", "Recipients in Cc line"), ("BCC - Expanded", "BCC – Expanded", "Recipients in Bcc line"), ): value = pst_export_row_value(row, *keys) if value: recipients.append(value) if not recipients: return None return ", ".join(recipients) def pst_export_manifest_recipients(tag_values: dict[str, str]) -> str | None: recipients = [ value for value in ( pst_export_normalized_text(tag_values.get("#To")), pst_export_normalized_text(tag_values.get("#CC")), pst_export_normalized_text(tag_values.get("#BCC")), ) if value ] if not recipients: return None return ", ".join(recipients) def pst_export_merge_message_metadata( existing: dict[str, object] | None, updates: dict[str, object], ) -> dict[str, object]: merged = dict(existing or {}) for key, value in updates.items(): if value in (None, "", [], {}): continue current = merged.get(key) if current in (None, "", [], {}): merged[key] = value return merged def parse_pst_export_results_csv( paths: list[Path], *, candidate_root: Path, ) -> dict[str, dict[str, dict[str, object]]]: metadata_by_pst_path: dict[str, dict[str, dict[str, object]]] = defaultdict(dict) for path in paths: headers, rows = load_normalized_csv_rows(path) header_set = set(headers) standard_format = PST_EXPORT_STANDARD_RESULTS_REQUIRED_HEADERS.issubset(header_set) archive_format = PST_EXPORT_ARCHIVE_RESULTS_REQUIRED_HEADERS.issubset(header_set) if not standard_format and not archive_format: continue for row in rows: source_item_id = normalize_pst_identifier( row.get("Export Item Id") if archive_format else row.get("Item Identity") ) if not source_item_id: continue resolved_pst_path, has_message_component = pst_export_resolve_pst_path( candidate_root, row.get("Export Item Path") if archive_format else row.get("Target Path"), ) if resolved_pst_path is None or not has_message_component: continue metadata_updates = { "author": pst_export_row_value(row, "Sender or Created by"), "compliance_tag": pst_export_row_value(row, "Compliance Tag"), "custodian": pst_export_row_value(row, "Location Name"), "date_created": normalize_datetime( pst_export_row_value(row, "Sent", "Received or Created") ), "date_modified": normalize_datetime(pst_export_row_value(row, "Modified Date")), "decode_status": pst_export_row_value(row, "Decode Status"), "document_path": pst_export_row_value(row, "Document Path"), "export_document_id": pst_export_row_value(row, "Document ID"), "export_item_id": pst_export_row_value(row, "Export Item Id", "ExportItem Id"), "export_item_path": pst_export_row_value(row, "Export Item Path"), "has_attachments": pst_export_row_value(row, "Has Attachments"), "item_identity": pst_export_row_value(row, "Item Identity"), "internet_message_id": pst_export_row_value(row, "Internet Message Id"), "location": pst_export_row_value(row, "Location"), "location_name": pst_export_row_value(row, "Location Name"), "modern_attachment_embedded_urls": pst_export_row_value(row, "Modern Attachment Embedded Urls"), "original_path": pst_export_row_value(row, "Original Path"), "preservation_original_url": pst_export_row_value(row, "Preservation Original Url"), "recipients": pst_export_combined_recipients(row), "retention_url": pst_export_row_value(row, "Retention Url"), "sidecar_source_item_id": source_item_id, "subject": normalize_generated_document_title(pst_export_row_value(row, "Subject or Title")), "target_path": pst_export_row_value(row, "Target Path"), "teams_metadata": pst_export_row_value(row, "Teams Metadata"), } pst_key = resolved_pst_path.as_posix() existing = metadata_by_pst_path[pst_key].get(source_item_id) metadata_by_pst_path[pst_key][source_item_id] = pst_export_merge_message_metadata(existing, metadata_updates) return {key: dict(value) for key, value in metadata_by_pst_path.items()} def parse_pst_export_manifest_xml( paths: list[Path], *, candidate_root: Path, ) -> dict[str, dict[str, dict[str, object]]]: metadata_by_pst_path: dict[str, dict[str, dict[str, object]]] = defaultdict(dict) for path in paths: try: root = parse_xml_document(path.read_bytes()) except (OSError, ET.ParseError): continue for document in root.iter(): if xml_local_name(document.tag) != "Document": continue source_item_id = normalize_pst_identifier(document.attrib.get("DocID")) if not source_item_id: continue tag_values: dict[str, str] = {} location_custodian = None location_uri = None for node in document.iter(): local_name = xml_local_name(node.tag) if local_name == "Tag": tag_name = normalize_inline_whitespace(str(node.attrib.get("TagName") or "")) if not tag_name: continue normalized_value = pst_export_normalized_text(node.attrib.get("TagValue")) if normalized_value: tag_values[tag_name] = normalized_value continue if local_name == "Custodian": location_custodian = pst_export_normalized_text(node.text) continue if local_name == "LocationURI": location_uri = pst_export_normalized_text(node.text) resolved_pst_path, has_message_component = pst_export_resolve_pst_path( candidate_root, tag_values.get("TargetPath"), ) if resolved_pst_path is None or not has_message_component: continue metadata_updates = { "author": pst_export_normalized_text(tag_values.get("#From")), "custodian": pst_export_normalized_text(tag_values.get("#Source")) or location_custodian, "date_created": normalize_datetime( tag_values.get("#DateSent") or tag_values.get("#DateReceived") or tag_values.get("#CreatedOn") ), "has_attachments": pst_export_normalized_text(tag_values.get("#HasAttachments")), "location": location_custodian or location_uri, "location_name": pst_export_normalized_text(tag_values.get("#Source")), "manifest_doc_id": source_item_id, "original_url": pst_export_normalized_text(tag_values.get("#OriginalUrl")), "recipients": pst_export_manifest_recipients(tag_values), "sidecar_source_item_id": source_item_id, "subject": normalize_generated_document_title(tag_values.get("#Subject")), "target_path": pst_export_normalized_text(tag_values.get("TargetPath")), } pst_key = resolved_pst_path.as_posix() existing = metadata_by_pst_path[pst_key].get(source_item_id) metadata_by_pst_path[pst_key][source_item_id] = pst_export_merge_message_metadata(existing, metadata_updates) return {key: dict(value) for key, value in metadata_by_pst_path.items()} def build_pst_export_message_match_records( message_metadata_by_pst_path: dict[str, dict[str, dict[str, object]]], ) -> dict[str, list[dict[str, object]]]: records_by_pst_path: dict[str, list[dict[str, object]]] = {} for pst_path, metadata_by_source_item in message_metadata_by_pst_path.items(): records: list[dict[str, object]] = [] for sidecar_source_item_id, message_metadata in sorted(metadata_by_source_item.items()): exact_match_ids: list[str] = [] seen_exact_match_ids: set[str] = set() for candidate in ( sidecar_source_item_id, message_metadata.get("sidecar_source_item_id"), message_metadata.get("item_identity"), message_metadata.get("manifest_doc_id"), message_metadata.get("export_item_id"), message_metadata.get("export_document_id"), ): normalized_candidate = normalize_pst_identifier(candidate) or pst_export_normalized_text(candidate) if not normalized_candidate or normalized_candidate in seen_exact_match_ids: continue seen_exact_match_ids.add(normalized_candidate) exact_match_ids.append(normalized_candidate) folder_match_paths: list[tuple[str, ...]] = [] seen_folder_match_paths: set[tuple[str, ...]] = set() for raw_path in ( message_metadata.get("target_path"), message_metadata.get("original_path"), message_metadata.get("document_path"), message_metadata.get("export_item_path"), ): folder_match_parts = pst_export_folder_match_parts(raw_path) if not folder_match_parts or folder_match_parts in seen_folder_match_paths: continue seen_folder_match_paths.add(folder_match_parts) folder_match_paths.append(folder_match_parts) records.append( { "date_created": normalize_datetime(message_metadata.get("date_created")), "exact_match_ids": exact_match_ids, "folder_match_paths": folder_match_paths, "internet_message_id": normalize_email_message_id(message_metadata.get("internet_message_id")), "message_metadata": dict(message_metadata), "subject_key": pst_export_match_text_key(message_metadata.get("subject")), } ) records_by_pst_path[pst_path] = records return records_by_pst_path def pst_export_message_search_text(message_metadata: dict[str, object] | None) -> str | None: if not message_metadata: return None lines = ["PST export metadata"] for label, key in ( ("Sidecar source item ID", "sidecar_source_item_id"), ("Export document ID", "export_document_id"), ("Export item ID", "export_item_id"), ("Item identity", "item_identity"), ("Location name", "location_name"), ("Location", "location"), ("Export target path", "target_path"), ("Export item path", "export_item_path"), ("Document path", "document_path"), ("Original path", "original_path"), ("Original URL", "original_url"), ("Preservation original URL", "preservation_original_url"), ("Retention URL", "retention_url"), ("Internet message ID", "internet_message_id"), ("Modern attachment embedded URLs", "modern_attachment_embedded_urls"), ("Teams metadata", "teams_metadata"), ("Compliance tag", "compliance_tag"), ("Decode status", "decode_status"), ): value = pst_export_normalized_text(message_metadata.get(key)) if value: lines.append(f"{label}: {value}") normalized = normalize_whitespace("\n".join(lines)) return normalized or None def pst_export_sidecar_custodian_candidate(message_metadata: dict[str, object] | None) -> str | None: if not message_metadata: return None for key in ("custodian", "location_name"): candidate = pst_export_clean_custodian_candidate(message_metadata.get(key)) if candidate: return candidate return None def pst_export_custodian_entity_hint( message_metadata: dict[str, object] | None, *, identifier_scope: str | None = None, ) -> dict[str, object] | None: custodian = pst_export_sidecar_custodian_candidate(message_metadata) location = pst_export_normalized_text(dict(message_metadata or {}).get("location")) if not custodian or not location: return None if normalize_entity_lookup_text(custodian) == normalize_entity_lookup_text(location): return None identifier: dict[str, object] = { "identifier_type": "external_id", "identifier_name": "pst_location", "display_value": location, "normalized_value": normalize_entity_lookup_text(location), "is_verified": 1, } normalized_scope = normalize_entity_lookup_text(identifier_scope or "") if normalized_scope: identifier["identifier_scope"] = normalized_scope return { "display_value": custodian, "identifiers": [identifier], } def merge_entity_hint_payload( existing_hints: object, *, role: str, hint: dict[str, object] | None, ) -> dict[str, object]: hints = dict(existing_hints) if isinstance(existing_hints, dict) else {} if hint is None: return hints role_hints = [ item for item in list(hints.get(role) or []) if isinstance(item, dict) ] hint_key = normalize_entity_lookup_text(hint.get("display_value") or "") for identifier in list(hint.get("identifiers") or []): if isinstance(identifier, dict): hint_key = entity_candidate_identifier_key(identifier) break if hint_key: for item in role_hints: item_key = normalize_entity_lookup_text(item.get("display_value") or "") for identifier in list(item.get("identifiers") or []): if isinstance(identifier, dict): item_key = entity_candidate_identifier_key(identifier) break if item_key == hint_key: return hints role_hints.append(hint) hints[role] = role_hints return hints def select_pst_export_message_metadata( normalized_message: dict[str, object], *, exact_metadata_by_source_item: dict[str, dict[str, object]] | None = None, message_match_records: list[dict[str, object]] | None = None, ) -> dict[str, object] | None: source_item_id = normalize_pst_identifier(normalized_message.get("source_item_id")) if source_item_id: exact_match = dict(exact_metadata_by_source_item or {}).get(source_item_id) if exact_match: return dict(exact_match) records = list(message_match_records or []) if not records: return None extracted = dict(normalized_message.get("extracted") or {}) email_threading = dict(extracted.get("email_threading") or {}) message_id = normalize_email_message_id(email_threading.get("message_id")) folder_match_parts = pst_export_folder_match_parts(normalized_message.get("source_folder_path")) date_created = normalize_datetime(extracted.get("date_created")) subject_key = pst_export_match_text_key(extracted.get("subject") or extracted.get("title")) candidates = records if source_item_id: exact_id_matches = [ record for record in candidates if source_item_id in list(record.get("exact_match_ids") or []) ] if len(exact_id_matches) == 1: return dict(exact_id_matches[0]["message_metadata"]) if exact_id_matches: candidates = exact_id_matches if message_id: message_id_matches = [ record for record in candidates if normalize_email_message_id(record.get("internet_message_id")) == message_id ] if len(message_id_matches) == 1: return dict(message_id_matches[0]["message_metadata"]) if message_id_matches: candidates = message_id_matches if folder_match_parts: folder_matches = [ record for record in candidates if any( pst_export_folder_parts_match(folder_match_parts, tuple(candidate_parts)) for candidate_parts in list(record.get("folder_match_paths") or []) ) ] if len(folder_matches) == 1: return dict(folder_matches[0]["message_metadata"]) if folder_matches: candidates = folder_matches if date_created: date_matches = [ record for record in candidates if normalize_datetime(record.get("date_created")) == date_created ] if len(date_matches) == 1: return dict(date_matches[0]["message_metadata"]) if date_matches: candidates = date_matches if subject_key: subject_matches = [ record for record in candidates if pst_export_match_text_key(record.get("subject_key")) == subject_key or pst_export_match_text_key(dict(record.get("message_metadata") or {}).get("subject")) == subject_key ] if len(subject_matches) == 1: return dict(subject_matches[0]["message_metadata"]) if subject_matches: candidates = subject_matches if len(candidates) == 1: return dict(candidates[0]["message_metadata"]) return None def apply_pst_export_message_metadata( extracted: dict[str, object], *, message_metadata: dict[str, object] | None, identifier_scope: str | None = None, ) -> dict[str, object]: if not message_metadata: return dict(extracted) enriched = dict(extracted) if not normalize_whitespace(str(enriched.get("date_created") or "")) and message_metadata.get("date_created"): enriched["date_created"] = message_metadata.get("date_created") if not normalize_whitespace(str(enriched.get("date_modified") or "")) and message_metadata.get("date_modified"): enriched["date_modified"] = message_metadata.get("date_modified") subject = normalize_generated_document_title(message_metadata.get("subject")) if subject and not normalize_whitespace(str(enriched.get("title") or "")): enriched["title"] = subject if subject and not normalize_whitespace(str(enriched.get("subject") or "")): enriched["subject"] = subject author = pst_export_normalized_text(message_metadata.get("author")) if author and not normalize_whitespace(str(enriched.get("author") or "")): enriched["author"] = author recipients = pst_export_normalized_text(message_metadata.get("recipients")) if recipients and not normalize_whitespace(str(enriched.get("recipients") or "")): enriched["recipients"] = recipients custodian_candidate = pst_export_sidecar_custodian_candidate(message_metadata) current_custodian = normalize_whitespace(str(enriched.get("custodian") or "")) if custodian_candidate and (not current_custodian or ("@" in custodian_candidate and "@" not in current_custodian)): enriched["custodian"] = custodian_candidate custodian_entity_hint = pst_export_custodian_entity_hint( message_metadata, identifier_scope=identifier_scope, ) if custodian_entity_hint is not None: enriched["entity_hints"] = merge_entity_hint_payload( enriched.get("entity_hints"), role="custodian", hint=custodian_entity_hint, ) if normalize_whitespace(str(enriched.get("content_type") or "")).lower() == "email": email_threading = dict(enriched.get("email_threading") or {}) if not normalize_email_message_id(email_threading.get("message_id")): internet_message_id = pst_export_normalized_text(message_metadata.get("internet_message_id")) if internet_message_id: email_threading["message_id"] = internet_message_id enriched["email_threading"] = email_threading search_text = pst_export_message_search_text(message_metadata) return append_extracted_search_context(enriched, [search_text] if search_text else []) def pst_export_enriched_message_file_hash( base_hash: object, *, message_metadata: dict[str, object] | None, ) -> str: return sha256_json_value( { "base_hash": normalize_whitespace(str(base_hash or "")) or None, "message_metadata": message_metadata or {}, } ) def detect_pst_export_root(candidate_root: Path) -> dict[str, object] | None: if not candidate_root.is_dir() or ".retriever" in candidate_root.parts: return None try: pst_paths = sorted( path for path in candidate_root.rglob("*.pst") if ".retriever" not in path.parts ) csv_paths = sorted( path for path in candidate_root.rglob("*.csv") if ".retriever" not in path.parts ) xml_paths = sorted( path for path in candidate_root.rglob("*.xml") if ".retriever" not in path.parts ) log_paths = sorted( path for path in candidate_root.rglob("*.log") if ".retriever" not in path.parts ) except OSError: return None if not pst_paths: return None results_csv_paths = [path for path in csv_paths if pst_export_results_csv_valid(path)] summary_csv_paths = [path for path in csv_paths if pst_export_summary_csv_valid(path)] manifest_paths = [path for path in xml_paths if pst_export_manifest_xml_valid(path)] trace_log_paths = [path for path in log_paths if pst_export_trace_log_valid(path)] if not (results_csv_paths or summary_csv_paths or manifest_paths or trace_log_paths): return None message_metadata_by_pst_path = parse_pst_export_results_csv(results_csv_paths, candidate_root=candidate_root) manifest_metadata_by_pst_path = parse_pst_export_manifest_xml(manifest_paths, candidate_root=candidate_root) for pst_path, records in manifest_metadata_by_pst_path.items(): target_records = message_metadata_by_pst_path.setdefault(pst_path, {}) for source_item_id, metadata in records.items(): target_records[source_item_id] = pst_export_merge_message_metadata( target_records.get(source_item_id), metadata, ) message_match_records_by_pst_path = build_pst_export_message_match_records(message_metadata_by_pst_path) owned_paths = { *(path.resolve() for path in results_csv_paths), *(path.resolve() for path in summary_csv_paths), *(path.resolve() for path in manifest_paths), *(path.resolve() for path in trace_log_paths), } message_sidecar_paths = [*results_csv_paths, *manifest_paths] message_sidecar_hash = ( sha256_json_value( { "files": { path.resolve().as_posix(): sha256_file(path) for path in sorted(message_sidecar_paths, key=lambda item: item.resolve().as_posix()) } } ) if message_sidecar_paths else None ) return { "message_metadata_by_pst_path": message_metadata_by_pst_path, "message_match_records_by_pst_path": message_match_records_by_pst_path, "message_sidecar_hash": message_sidecar_hash, "owned_paths": owned_paths, "pst_paths": [path.resolve() for path in pst_paths], "root": candidate_root, } def find_pst_export_roots( root: Path, recursive: bool, ) -> list[dict[str, object]]: candidates: set[Path] = set() if recursive: try: for csv_path in root.rglob("*.csv"): if ".retriever" in csv_path.parts: continue if pst_export_results_csv_valid(csv_path) or pst_export_summary_csv_valid(csv_path): candidates.add(csv_path.parent) for xml_path in root.rglob("*.xml"): if ".retriever" in xml_path.parts: continue if pst_export_manifest_xml_valid(xml_path): candidates.add(xml_path.parent) for log_path in root.rglob("*.log"): if ".retriever" in log_path.parts: continue if pst_export_trace_log_valid(log_path): candidates.add(log_path.parent) except OSError: return [] else: candidates.add(root) descriptors: list[dict[str, object]] = [] accepted_roots: list[Path] = [] for candidate in sorted(candidates, key=lambda path: (len(path.parts), path.as_posix())): if any(parent == candidate or parent in candidate.parents for parent in accepted_roots): continue descriptor = detect_pst_export_root(candidate) if descriptor is None: continue descriptors.append(descriptor) accepted_roots.append(candidate) return descriptors def detect_gmail_export_root(candidate_root: Path) -> dict[str, object] | None: if not candidate_root.is_dir() or ".retriever" in candidate_root.parts: return None archive_browser_path = candidate_root / GMAIL_EXPORT_ARCHIVE_BROWSER_FILE try: mbox_paths = sorted( path for path in candidate_root.rglob("*.mbox") if ".retriever" not in path.parts ) except OSError: return None if not mbox_paths: return None try: metadata_csv_paths = sorted(path for path in candidate_root.rglob("*.csv") if gmail_metadata_csv_valid(path)) drive_links_paths = sorted(path for path in candidate_root.rglob("*.csv") if gmail_drive_links_csv_valid(path)) drive_export_dirs = sorted( path for path in candidate_root.rglob("*") if path.is_dir() and GMAIL_EXPORT_DRIVE_FOLDER_PATTERN.search(path.name) ) drive_export_metadata_paths = sorted( path for path in candidate_root.rglob("*-metadata.xml") if "Drive_Link_Export" in path.name ) drive_export_error_paths = sorted( path for path in candidate_root.rglob("*-errors.csv") if "Drive_Link_Export" in path.name ) auxiliary_metadata_xml_paths = sorted( path for path in candidate_root.rglob("*-metadata.xml") if "Drive_Link_Export" not in path.name ) auxiliary_error_xml_paths = sorted(path for path in candidate_root.rglob("*-errors.xml")) result_count_paths = sorted(path for path in candidate_root.rglob("*-result-counts.csv")) valid_result_count_paths = [path for path in result_count_paths if gmail_result_counts_csv_valid(path)] md5_paths = sorted(path for path in candidate_root.rglob("*.md5")) zip_paths = sorted(path for path in candidate_root.rglob("*.zip")) except OSError: return None full_export_detected = bool( archive_browser_path.exists() or metadata_csv_paths or drive_links_paths or drive_export_dirs or drive_export_metadata_paths ) mbox_family_keys = { key for path in mbox_paths for key in gmail_export_family_keys(path) } degraded_sidecar_paths = [ *auxiliary_error_xml_paths, *valid_result_count_paths, *md5_paths, ] degraded_export_detected = any( gmail_export_family_keys(path) & mbox_family_keys for path in degraded_sidecar_paths ) if not (full_export_detected or degraded_export_detected): return None owned_zip_paths = sorted( path for path in zip_paths if gmail_export_family_keys(path) & mbox_family_keys ) email_metadata_by_message_id = parse_gmail_metadata_csv(metadata_csv_paths) drive_links_by_message_id = parse_gmail_drive_links_csv(drive_links_paths) drive_export_files = gmail_drive_export_files(drive_export_dirs) file_paths_by_name = {path.name: path for path in drive_export_files} drive_records_by_item_id: dict[str, dict[str, object]] = {} drive_documents: list[dict[str, object]] = [] for export_path in drive_export_files: drive_item_id = gmail_normalized_drive_item_id(gmail_drive_item_id_from_export_file_name(export_path)) record = { "author": None, "collaborators": [], "date_created": None, "date_modified": None, "document_type": None, "drive_item_id": drive_item_id, "error": None, "file_name": export_path.name, "file_path": export_path, "linked_message_ids": [], "linked_subjects": [], "others": [], "source_hash": None, "title": normalize_generated_document_title(export_path.stem.rsplit("_", 1)[0]) or None, "viewers": [], } drive_documents.append(record) if drive_item_id: drive_records_by_item_id[drive_item_id] = record parsed_drive_metadata = parse_gmail_drive_export_metadata( drive_export_metadata_paths, file_paths_by_name=file_paths_by_name, ) for drive_item_id, metadata in parsed_drive_metadata.items(): record = drive_records_by_item_id.get(drive_item_id) if record is None: record = { "author": None, "collaborators": [], "date_created": None, "date_modified": None, "document_type": None, "drive_item_id": drive_item_id, "error": None, "file_name": metadata.get("file_name"), "file_path": metadata.get("file_path"), "linked_message_ids": [], "linked_subjects": [], "others": [], "source_hash": None, "title": None, "viewers": [], } drive_records_by_item_id[drive_item_id] = record if record.get("file_path") is not None: drive_documents.append(record) record.update( { "author": metadata.get("author"), "collaborators": list(metadata.get("collaborators") or []), "date_created": metadata.get("date_created"), "date_modified": metadata.get("date_modified"), "document_type": metadata.get("document_type"), "file_name": metadata.get("file_name"), "file_path": metadata.get("file_path"), "others": list(metadata.get("others") or []), "source_hash": metadata.get("source_hash"), "title": metadata.get("title"), "viewers": list(metadata.get("viewers") or []), } ) drive_errors_by_item_id = parse_gmail_drive_export_errors(drive_export_error_paths) linked_drive_records_by_message_id: dict[str, list[dict[str, object]]] = defaultdict(list) linked_drive_attachment_records_by_message_id: dict[str, list[dict[str, object]]] = defaultdict(list) for message_id, link_rows in drive_links_by_message_id.items(): seen_drive_items: set[str] = set() for link_row in link_rows: drive_item_id = gmail_normalized_drive_item_id(link_row.get("drive_item_id")) if not drive_item_id or drive_item_id in seen_drive_items: continue seen_drive_items.add(drive_item_id) drive_record = drive_records_by_item_id.get(drive_item_id) summary = gmail_drive_document_link_summary( { **dict(drive_record or {}), "drive_item_id": drive_item_id, "error": drive_errors_by_item_id.get(drive_item_id), } ) linked_drive_records_by_message_id[message_id].append(summary) if drive_record is not None: if message_id not in drive_record["linked_message_ids"]: drive_record["linked_message_ids"].append(message_id) subject = normalize_generated_document_title( (email_metadata_by_message_id.get(message_id) or {}).get("subject") ) if subject and subject not in drive_record["linked_subjects"]: drive_record["linked_subjects"].append(subject) drive_record["error"] = drive_errors_by_item_id.get(drive_item_id) if isinstance(drive_record.get("file_path"), Path): linked_drive_attachment_records_by_message_id[message_id].append(drive_record) deduplicated_drive_documents: dict[str, dict[str, object]] = {} for record in drive_documents: file_path = record.get("file_path") drive_item_id = gmail_normalized_drive_item_id(record.get("drive_item_id")) if isinstance(file_path, Path): key = f"path:{file_path.resolve().as_posix()}" elif drive_item_id: key = f"drive:{drive_item_id}" else: key = f"file:{normalize_whitespace(str(record.get('file_name') or ''))}" current = deduplicated_drive_documents.get(key) if current is None or gmail_drive_record_preference_key(record) > gmail_drive_record_preference_key(current): deduplicated_drive_documents[key] = record drive_documents = sorted( deduplicated_drive_documents.values(), key=lambda record: ( str(record.get("file_path") or ""), str(record.get("file_name") or ""), str(record.get("drive_item_id") or ""), ), ) owned_paths = { *(path.resolve() for path in mbox_paths), *(path.resolve() for path in metadata_csv_paths), *(path.resolve() for path in drive_links_paths), *(path.resolve() for path in auxiliary_metadata_xml_paths), *(path.resolve() for path in auxiliary_error_xml_paths), *(path.resolve() for path in drive_export_metadata_paths), *(path.resolve() for path in drive_export_error_paths), *(path.resolve() for path in result_count_paths), *(path.resolve() for path in md5_paths), *(path.resolve() for path in owned_zip_paths), *(path.resolve() for path in drive_export_files), } if archive_browser_path.exists(): owned_paths.add(archive_browser_path.resolve()) message_sidecar_paths = [ *metadata_csv_paths, *drive_links_paths, *drive_export_metadata_paths, *drive_export_error_paths, *drive_export_files, ] message_sidecar_hash = sha256_json_value( { "files": { path.resolve().as_posix(): sha256_file(path) for path in sorted(message_sidecar_paths, key=lambda item: item.resolve().as_posix()) } } ) return { "drive_documents": drive_documents, "drive_records_by_item_id": drive_records_by_item_id, "email_metadata_by_message_id": email_metadata_by_message_id, "linked_drive_attachment_records_by_message_id": linked_drive_attachment_records_by_message_id, "linked_drive_records_by_message_id": linked_drive_records_by_message_id, "mbox_paths": mbox_paths, "message_sidecar_hash": message_sidecar_hash, "owned_paths": owned_paths, "root": candidate_root, } def find_gmail_export_roots( root: Path, recursive: bool, allowed_file_types: set[str] | None, ) -> list[dict[str, object]]: candidates: set[Path] = set() if recursive: try: for archive_browser_path in root.rglob(GMAIL_EXPORT_ARCHIVE_BROWSER_FILE): if ".retriever" in archive_browser_path.parts: continue candidates.add(archive_browser_path.parent) for metadata_path in root.rglob("*-metadata.csv"): if ".retriever" in metadata_path.parts: continue candidates.add(metadata_path.parent) for drive_links_path in root.rglob("*-drive-links.csv"): if ".retriever" in drive_links_path.parts: continue candidates.add(drive_links_path.parent) for result_count_path in root.rglob("*-result-counts.csv"): if ".retriever" in result_count_path.parts: continue if gmail_result_counts_csv_valid(result_count_path): candidates.add(result_count_path.parent) for error_path in root.rglob("*-errors.xml"): if ".retriever" in error_path.parts: continue candidates.add(error_path.parent) for md5_path in root.rglob("*.md5"): if ".retriever" in md5_path.parts: continue candidates.add(md5_path.parent) for export_dir in root.rglob("*"): if not export_dir.is_dir() or ".retriever" in export_dir.parts: continue if GMAIL_EXPORT_DRIVE_FOLDER_PATTERN.search(export_dir.name): candidates.add(export_dir.parent) except OSError: return [] else: candidates.add(root) descriptors: list[dict[str, object]] = [] accepted_roots: list[Path] = [] for candidate in sorted(candidates, key=lambda path: (len(path.parts), path.as_posix())): if any(parent == candidate or parent in candidate.parents for parent in accepted_roots): continue descriptor = detect_gmail_export_root(candidate) if descriptor is None: continue descriptors.append(descriptor) accepted_roots.append(candidate) return descriptors file_type = normalize_extension(path) if file_type not in SUPPORTED_FILE_TYPES: raise RetrieverError(f"Unsupported file type: .{file_type or '(none)'}") if file_type in TEXT_FILE_TYPES: return extract_plain_text_file(path) if file_type in IMAGE_NATIVE_PREVIEW_FILE_TYPES: return extract_native_preview_only_file(path, explicit_content_type="Image") if file_type == "rtf": return extract_rtf_file(path) if file_type == "pdf": return extract_pdf_file(path) if file_type == "docx": return extract_docx_file(path) if file_type == "pptx": return extract_pptx_file(path) if file_type == "xls": return extract_xls_file(path) if file_type == "xlsx": return extract_xlsx_file(path) if file_type == "eml": return extract_eml_file(path, include_attachments=include_attachments) if file_type == "msg": return extract_msg_file(path, include_attachments=include_attachments) raise RetrieverError(f"Unsupported file type: .{file_type}") def relative_document_path(root: Path, path: Path) -> str: return path.resolve().relative_to(root.resolve()).as_posix() def path_is_at_or_under(path: Path, parent: Path) -> bool: resolved_path = path.resolve() resolved_parent = parent.resolve() return resolved_path == resolved_parent or resolved_parent in resolved_path.parents def build_ingest_scan_scope(root: Path, raw_paths: list[str] | None) -> dict[str, object]: root = root.resolve() if not raw_paths: return { "is_full_workspace": True, "paths": [root], "dir_prefixes": [], "exact_rel_paths": set(), "display_paths": [], } scan_paths: list[Path] = [] dir_prefixes: set[str] = set() exact_rel_paths: set[str] = set() display_paths: list[str] = [] for raw_path in raw_paths: raw_text = normalize_whitespace(str(raw_path or "")) if not raw_text: raise RetrieverError("Ingest --path cannot be blank.") candidate = Path(raw_text).expanduser() candidate = candidate if candidate.is_absolute() else root / candidate if not candidate.exists(): raise RetrieverError(f"Ingest path does not exist: {raw_text}") resolved_candidate = candidate.resolve() if not path_is_at_or_under(resolved_candidate, root): raise RetrieverError(f"Ingest path must be inside the workspace root: {raw_text}") rel_path = resolved_candidate.relative_to(root).as_posix() rel_parts = resolved_candidate.relative_to(root).parts if ".retriever" in rel_parts: raise RetrieverError("Ingest --path cannot target the .retriever state directory.") if rel_path in {"", "."}: return { "is_full_workspace": True, "paths": [root], "dir_prefixes": [], "exact_rel_paths": set(), "display_paths": [], } scan_paths.append(resolved_candidate) display_paths.append(rel_path) if resolved_candidate.is_dir(): dir_prefixes.add(rel_path.rstrip("/") + "/") else: exact_rel_paths.add(rel_path) deduped_scan_paths: list[Path] = [] for scan_path in sorted(set(scan_paths), key=lambda item: (len(item.parts), item.as_posix())): if any(parent == scan_path or parent in scan_path.parents for parent in deduped_scan_paths if parent.is_dir()): continue deduped_scan_paths.append(scan_path) return { "is_full_workspace": False, "paths": deduped_scan_paths, "dir_prefixes": sorted(dir_prefixes), "exact_rel_paths": exact_rel_paths, "display_paths": sorted(set(display_paths)), } def ingest_scan_scope_contains_rel_path(scan_scope: dict[str, object] | None, rel_path: str | None) -> bool: if scan_scope is None or bool(scan_scope.get("is_full_workspace")): return True normalized = normalize_whitespace(str(rel_path or "")).replace("\\", "/").lstrip("/") if not normalized: return False if normalized in set(scan_scope.get("exact_rel_paths") or set()): return True return any(normalized.startswith(str(prefix)) for prefix in list(scan_scope.get("dir_prefixes") or [])) def collect_files( root: Path, recursive: bool, allowed_file_types: set[str] | None, scan_scope: dict[str, object] | None = None, ) -> list[Path]: scan_paths = list((scan_scope or {}).get("paths") or [root]) files: set[Path] = set() for scan_path in scan_paths: if scan_path.is_file(): candidates = [scan_path] else: candidates = scan_path.rglob("*") if recursive else scan_path.iterdir() for path in candidates: if not path_is_at_or_under(path, root): continue if not ingest_scan_scope_contains_rel_path(scan_scope, relative_document_path(root, path)): continue if path.is_dir(): if path.name == ".retriever": continue if not recursive: continue continue if ".retriever" in path.resolve().relative_to(root.resolve()).parts: continue file_type = normalize_extension(path) if allowed_file_types and file_type not in allowed_file_types: continue if not file_type: continue files.add(path) return sorted(files) def collect_ingest_scope_directories(scan_scope: dict[str, object]) -> list[Path]: return [Path(path) for path in list(scan_scope.get("paths") or []) if Path(path).is_dir()] def dedupe_source_descriptors_by_root(descriptors: list[dict[str, object]]) -> list[dict[str, object]]: descriptors_by_root: dict[str, dict[str, object]] = {} for descriptor in descriptors: root = descriptor.get("root") if root is None: continue descriptors_by_root[Path(root).resolve().as_posix()] = descriptor return [descriptors_by_root[key] for key in sorted(descriptors_by_root)] def find_scoped_source_roots( finder, root: Path, recursive: bool, scan_scope: dict[str, object], *args, ) -> list[dict[str, object]]: if bool(scan_scope.get("is_full_workspace")): return finder(root, recursive, *args) descriptors: list[dict[str, object]] = [] for scan_dir in collect_ingest_scope_directories(scan_scope): descriptors.extend(finder(scan_dir, recursive, *args)) return dedupe_source_descriptors_by_root(descriptors) def collect_production_dat_paths(root: Path, recursive: bool, scan_scope: dict[str, object]) -> list[Path]: if bool(scan_scope.get("is_full_workspace")): dat_paths = list(root.rglob("*.dat")) if recursive else list(root.glob("*.dat")) if not recursive: for child in root.iterdir(): if child.is_dir() and child.name != ".retriever": data_dir = child / "DATA" if data_dir.is_dir(): dat_paths.extend(sorted(data_dir.glob("*.dat"))) return dat_paths dat_paths: list[Path] = [] for scan_dir in collect_ingest_scope_directories(scan_scope): if recursive: dat_paths.extend(scan_dir.rglob("*.dat")) else: dat_paths.extend(scan_dir.glob("*.dat")) for child in scan_dir.iterdir(): if child.is_dir() and child.name != ".retriever": data_dir = child / "DATA" if data_dir.is_dir(): dat_paths.extend(sorted(data_dir.glob("*.dat"))) return dat_paths def collect_scoped_existing_production_rows( workspace_root: Path, rows: list[sqlite3.Row], scan_scope: dict[str, object], ) -> list[sqlite3.Row]: if bool(scan_scope.get("is_full_workspace")): return rows scoped_rows: list[sqlite3.Row] = [] for row in rows: rel_root = str(row["rel_root"]) candidate_root = workspace_root / rel_root if ingest_scan_scope_contains_rel_path(scan_scope, rel_root) or any( path_is_at_or_under(Path(path), candidate_root) for path in list(scan_scope.get("paths") or []) ): scoped_rows.append(row) return scoped_rows def normalize_production_header(raw_header: str) -> str: normalized = re.sub(r"[^a-z0-9]+", " ", raw_header.strip().lower()).strip() squashed = normalized.replace(" ", "") return PRODUCTION_DAT_HEADER_ALIASES.get(normalized) or PRODUCTION_DAT_HEADER_ALIASES.get(squashed) or normalized.replace(" ", "_") def parse_concordance_rows(raw_bytes: bytes) -> list[list[str]]: text = raw_bytes.decode("latin-1") reader = csv.reader(io.StringIO(text), delimiter="\x14", quotechar="\xfe") return [[field.strip().strip("\ufeff") for field in row] for row in reader if any(field.strip() for field in row)] def parse_generic_delimited_rows(raw_bytes: bytes) -> list[list[str]]: decoded, _, _ = decode_bytes(raw_bytes) sample = decoded[:4096] try: dialect = csv.Sniffer().sniff(sample, delimiters=",|\t;") except csv.Error: dialect = csv.excel reader = csv.reader(io.StringIO(decoded), dialect) return [[field.strip().strip("\ufeff") for field in row] for row in reader if any(field.strip() for field in row)] def parse_production_metadata_load(path: Path) -> dict[str, object]: raw_bytes = path.read_bytes() rows = parse_concordance_rows(raw_bytes) if b"\x14" in raw_bytes else parse_generic_delimited_rows(raw_bytes) if not rows: raise RetrieverError(f"Production metadata load file is empty: {path}") headers = [normalize_production_header(header) for header in rows[0]] if "begin_bates" not in headers or "end_bates" not in headers: raise RetrieverError(f"Production metadata load file missing Bates headers: {path}") records: list[dict[str, str]] = [] for values in rows[1:]: record: dict[str, str] = {} for index, header in enumerate(headers): if index >= len(values): continue value = values[index].strip() if value: record[header] = value if not record: continue if "text_path" not in record: for key, value in list(record.items()): if "text" in key and normalize_extension(Path(value)) == "txt": record["text_path"] = value break if "native_path" not in record: for key, value in list(record.items()): if ("file" in key or "native" in key or "path" in key) and normalize_extension(Path(value)) not in {"", "txt"}: record["native_path"] = value break if record.get("begin_bates"): records.append(record) return {"headers": headers, "rows": records} def parse_production_image_load(path: Path | None) -> list[dict[str, str]]: if path is None or not path.exists(): return [] decoded, _, _ = decode_bytes(path.read_bytes()) reader = csv.reader(io.StringIO(decoded)) rows: list[dict[str, str]] = [] for values in reader: if len(values) < 3: continue page_bates = values[0].strip() image_path = values[2].strip() if not page_bates or not image_path: continue rows.append( { "page_bates": page_bates, "volume_name": values[1].strip() if len(values) > 1 else "", "image_path": image_path, "is_first_page": values[3].strip().upper() == "Y" if len(values) > 3 else False, } ) return rows def find_production_load_file(candidate_root: Path, extension: str) -> Path | None: preferred_dir = candidate_root / "DATA" candidates: list[Path] = [] if preferred_dir.is_dir(): candidates.extend(sorted(preferred_dir.glob(f"*.{extension}"))) candidates.extend(sorted(path for path in candidate_root.glob(f"*.{extension}") if path not in candidates)) return candidates[0] if candidates else None def production_signature_for_root(workspace_root: Path, candidate_root: Path) -> dict[str, object] | None: workspace_root = workspace_root.resolve() candidate_root = candidate_root.resolve() metadata_load_path = find_production_load_file(candidate_root, "dat") if metadata_load_path is None: return None has_payload_dirs = any((candidate_root / name).is_dir() for name in ("TEXT", "IMAGES", "NATIVES")) if not has_payload_dirs: return None try: metadata = parse_production_metadata_load(metadata_load_path) except Exception: return None headers = set(metadata["headers"]) if not {"begin_bates", "end_bates"}.issubset(headers): return None image_load_path = find_production_load_file(candidate_root, "opt") return { "root": candidate_root, "rel_root": candidate_root.relative_to(workspace_root).as_posix(), "production_name": candidate_root.name, "metadata_load_path": metadata_load_path, "image_load_path": image_load_path, "source_type": "concordance-dat-opt" if image_load_path else "concordance-dat", } def find_production_root_signatures( workspace_root: Path, recursive: bool, connection: sqlite3.Connection | None = None, scan_scope: dict[str, object] | None = None, ) -> list[dict[str, object]]: scan_scope = scan_scope or build_ingest_scan_scope(workspace_root, None) signatures: dict[str, dict[str, object]] = {} if connection is not None and table_exists(connection, "productions"): rows = connection.execute("SELECT rel_root, production_name, metadata_load_rel_path, image_load_rel_path, source_type FROM productions").fetchall() for row in collect_scoped_existing_production_rows(workspace_root, rows, scan_scope): candidate_root = workspace_root / row["rel_root"] if candidate_root.exists(): signatures[row["rel_root"]] = { "root": candidate_root, "rel_root": row["rel_root"], "production_name": row["production_name"], "metadata_load_path": workspace_root / row["metadata_load_rel_path"], "image_load_path": (workspace_root / row["image_load_rel_path"]) if row["image_load_rel_path"] else None, "source_type": row["source_type"], } dat_paths = collect_production_dat_paths(workspace_root, recursive, scan_scope) for dat_path in dat_paths: if ".retriever" in dat_path.parts: continue candidate_root = dat_path.parent.parent if dat_path.parent.name.lower() == "data" else dat_path.parent if candidate_root == workspace_root / ".retriever": continue try: signature = production_signature_for_root(workspace_root, candidate_root) except Exception: signature = None if signature is not None: signatures[str(signature["rel_root"])] = signature return [signatures[key] for key in sorted(signatures)] def normalize_source_reference(raw_value: str) -> list[str]: normalized = raw_value.strip().strip('"').strip("'").replace("\\", "/") normalized = re.sub(r"^[.]/+", "", normalized) normalized = normalized.lstrip("/") return [part for part in normalized.split("/") if part and part != "."] def source_reference_suffixes(parts: list[str]) -> list[list[str]]: seen: set[tuple[str, ...]] = set() suffixes: list[list[str]] = [] for index in range(len(parts)): suffix = tuple(parts[index:]) if not suffix or suffix in seen: continue seen.add(suffix) suffixes.append(list(suffix)) return suffixes def resolve_case_insensitive_relative_path(base_dir: Path, parts: list[str]) -> Path | None: current = base_dir for part in parts: if not current.is_dir(): return None entries = {child.name.lower(): child for child in current.iterdir()} candidate = entries.get(part.lower()) if candidate is None: return None current = candidate return current def resolve_production_source_path(workspace_root: Path, production_root: Path, raw_value: str | None) -> Path | None: if raw_value is None: return None stripped = raw_value.strip().strip('"').strip("'") if not stripped: return None raw_path = Path(stripped) if raw_path.is_absolute(): try: absolute = raw_path.resolve() except Exception: absolute = raw_path return absolute if absolute.exists() else None parts = normalize_source_reference(stripped) suffixes = source_reference_suffixes(parts) if not suffixes: return None # Preserve existing literal resolution before we attempt any fallback prefix stripping. literal_parts = suffixes[0] resolved = resolve_case_insensitive_relative_path(production_root, literal_parts) if resolved is not None: return resolved resolved = resolve_case_insensitive_relative_path(workspace_root, literal_parts) if resolved is not None: return resolved for candidate_parts in suffixes[1:]: resolved = resolve_case_insensitive_relative_path(production_root, candidate_parts) if resolved is not None: return resolved for candidate_parts in suffixes[1:]: resolved = resolve_case_insensitive_relative_path(workspace_root, candidate_parts) if resolved is not None: return resolved return None def production_logical_rel_path(production_rel_root: str, control_number: str) -> str: production_slug = sanitize_storage_filename(Path(production_rel_root).name) control_slug = sanitize_storage_filename(control_number) return Path(INTERNAL_REL_PATH_PREFIX) / "productions" / production_slug / "documents" / f"{control_slug}.logical" def production_preview_page_asset_refs( rel_path: str, control_number: str, image_paths: list[Path], ) -> list[dict[str, object]]: preview_base = preview_base_path_for_rel_path(rel_path) page_dir_name = f"{sanitize_storage_filename(control_number)}-pages" page_suffix = preview_image_output_suffix() refs: list[dict[str, object]] = [] for index, image_path in enumerate(image_paths, start=1): page_file_name = f"page-{index:04d}{page_suffix}" rel_preview_path = preview_base / page_dir_name / page_file_name refs.append( { "ordinal": index, "label": f"Page {index}", "source_path": str(image_path), "rel_preview_path": rel_preview_path.as_posix(), "html_src": urllib_request.pathname2url(f"{page_dir_name}/{page_file_name}"), } ) return refs def production_preview_page_assets( preview_image_refs: list[dict[str, object]], *, max_dimension: int | None = None, ) -> list[dict[str, object]]: assets: list[dict[str, object]] = [] for ref in preview_image_refs: source_path = Path(str(ref.get("source_path") or "")) if not source_path.exists(): continue preview_raster = image_path_preview_raster(source_path, max_dimension=max_dimension) if preview_raster is None: continue preview_bytes, _, _ = preview_raster rel_preview_path = normalize_whitespace(str(ref.get("rel_preview_path") or "")) if not rel_preview_path: continue assets.append( { "ordinal": int(ref.get("ordinal") or 0), "label": str(ref.get("label") or ""), "rel_preview_path": rel_preview_path, "payload": preview_bytes, } ) return assets def infer_production_title(control_number: str, text_content: str, native_path: Path | None) -> str: for line in text_content.splitlines(): candidate = normalize_whitespace(line) if candidate: if re.match(r"^(From|To|Cc|Bcc|Sent|Date|Subject):\s*", candidate, flags=re.IGNORECASE): continue return normalize_generated_document_title(candidate[:220]) or candidate[:220] if native_path is not None: return normalize_generated_document_title(native_path.stem or native_path.name) or native_path.stem or native_path.name return control_number def build_production_preview_html( *, document_title: str, control_number: str, production_name: str, begin_bates: str, end_bates: str, begin_attachment: str | None, end_attachment: str | None, text_content: str, page_images: list[dict[str, object]], page_image_note: str | None = None, ) -> str: headers = { passive_field_label("production_name"): production_name, passive_field_label("control_number"): control_number, passive_field_label("begin_bates"): begin_bates, passive_field_label("end_bates"): end_bates, passive_field_label("begin_attachment"): begin_attachment or "", passive_field_label("end_attachment"): end_attachment or "", } sections: list[str] = [] if text_content.strip(): sections.append( "

Linked Text

"
            + html.escape(text_content)
            + "
" ) if page_images: image_sections: list[str] = ["

Produced Pages

"] if page_image_note: image_sections.append(f"

{html.escape(page_image_note)}

") for image in page_images: label = html.escape(str(image["label"])) src = html.escape(str(image["src"])) image_sections.append(f'
{label}
{label}
') image_sections.append("
") sections.append("".join(image_sections)) elif page_image_note: sections.append(f"

Produced Pages

{html.escape(page_image_note)}

") body_html = "".join(sections) or "

No linked text or page images were available for this production document.

" head_html = """ """.strip() return build_html_preview(headers, body_html=body_html, document_title=document_title, head_html=head_html) def attachment_preview_link_label(row: sqlite3.Row) -> str: file_name = normalize_whitespace(str(row["file_name"] or "")) title = normalize_generated_document_title(row["title"]) control_number = normalize_whitespace(str(row["control_number"] or "")) if file_name and not file_name.lower().endswith(".logical"): return file_name if title: return title if file_name: return file_name return control_number or "Attachment" def relative_preview_href(path: Path, parent_preview_path: Path, target_fragment: object = None) -> str: href = urllib_request.pathname2url( os.path.relpath(str(path), start=str(parent_preview_path.parent)) ) return append_preview_fragment(href, target_fragment) def parse_calendar_attachment_metadata(paths: dict[str, Path], child_row: sqlite3.Row) -> dict[str, object] | None: content_type = normalize_whitespace(str(child_row["content_type"] or "")) file_type = normalize_whitespace(str(child_row["file_type"] or "")).lower() if not file_type: file_type = normalize_extension(Path(str(child_row["file_name"] or child_row["rel_path"] or ""))) if content_type != "Calendar" and file_type not in ICALENDAR_FILE_TYPES: return None rel_path = normalize_whitespace(str(child_row["rel_path"] or "")) if not rel_path: return None attachment_path = document_absolute_path(paths, rel_path) if not attachment_path.exists(): return None try: decoded, _, _ = decode_bytes(attachment_path.read_bytes()) except OSError: return None metadata = parse_icalendar_event_metadata(decoded) if not metadata: return None if not any(metadata.get(key) for key in ("summary", "when", "organizer", "attendees_display", "conference_url")): return None return metadata def build_calendar_attachment_preview_link( *, paths: dict[str, Path], child_row: sqlite3.Row, child_preview_path: Path, parent_preview_path: Path, ) -> dict[str, str] | None: metadata = parse_calendar_attachment_metadata(paths, child_row) if metadata is None: return None title = ( normalize_generated_document_title(metadata.get("summary")) or normalize_generated_document_title(child_row["title"]) or attachment_preview_link_label(child_row) ) organizer = metadata.get("organizer") or normalize_whitespace(str(child_row["author"] or "")) or None attendees = metadata.get("attendees_display") or normalize_whitespace(str(child_row["recipients"] or "")) or None return { "kind": "calendar_invite", "href": relative_preview_href(child_preview_path, parent_preview_path), "label": attachment_preview_link_label(child_row), "detail": normalize_whitespace(str(child_row["control_number"] or "")), "title": str(title or "Calendar invite"), "when": str(metadata.get("when") or ""), "organizer": str(organizer or ""), "attendees": str(attendees or ""), "location": str(metadata.get("location") or ""), "join_href": str(metadata.get("conference_url") or ""), "status": str(summarize_icalendar_invite_status(metadata) or ""), } def build_document_attachment_preview_links( paths: dict[str, Path], connection: sqlite3.Connection, parent_preview_path: Path, child_rows: list[sqlite3.Row], ) -> list[dict[str, str]]: links: list[dict[str, str]] = [] for child_row in child_rows: child_preview_target = default_preview_target(paths, child_row, connection) child_preview_abs_path = str( child_preview_target.get("file_abs_path") or child_preview_target.get("abs_path") or "" ).split("#", 1)[0] child_preview_path = Path(child_preview_abs_path) if not child_preview_path.exists(): continue calendar_link = build_calendar_attachment_preview_link( paths=paths, child_row=child_row, child_preview_path=child_preview_path, parent_preview_path=parent_preview_path, ) if calendar_link is not None: links.append(calendar_link) continue detail = normalize_whitespace(str(child_row["control_number"] or "")) links.append( { "href": relative_preview_href(child_preview_path, parent_preview_path), "label": attachment_preview_link_label(child_row), "detail": detail, } ) return links def conversation_attachment_links_by_document_id( connection: sqlite3.Connection, paths: dict[str, Path], *, segment_preview_path: Path, documents: list[dict[str, object]], ) -> dict[int, list[dict[str, str]]]: parent_ids = [ int(document["id"]) for document in documents if document.get("parent_document_id") is None ] if not parent_ids: return {} child_rows = connection.execute( f""" SELECT * FROM documents WHERE parent_document_id IN ({", ".join("?" for _ in parent_ids)}) AND lifecycle_status NOT IN ('missing', 'deleted') AND COALESCE(child_document_kind, ?) = ? ORDER BY parent_document_id ASC, CASE WHEN control_number_attachment_sequence IS NULL THEN 1 ELSE 0 END ASC, control_number_attachment_sequence ASC, control_number ASC, id ASC """, (*parent_ids, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT), ).fetchall() child_rows_by_parent_id: dict[int, list[sqlite3.Row]] = defaultdict(list) for child_row in child_rows: child_rows_by_parent_id[int(child_row["parent_document_id"])].append(child_row) return { parent_id: build_document_attachment_preview_links( paths, connection, segment_preview_path, rows, ) for parent_id, rows in child_rows_by_parent_id.items() } def sync_document_attachment_preview_links( connection: sqlite3.Connection, paths: dict[str, Path], document_id: int, ) -> int: preview_rows = connection.execute( """ SELECT rel_preview_path FROM document_previews WHERE document_id = ? AND preview_type = 'html' ORDER BY ordinal ASC, id ASC """, (document_id,), ).fetchall() if not preview_rows: return 0 child_rows = connection.execute( """ SELECT * FROM documents WHERE parent_document_id = ? AND lifecycle_status NOT IN ('missing', 'deleted') AND COALESCE(child_document_kind, ?) = ? ORDER BY CASE WHEN control_number_attachment_sequence IS NULL THEN 1 ELSE 0 END ASC, control_number_attachment_sequence ASC, control_number ASC, id ASC """, (document_id, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT), ).fetchall() updated = 0 for preview_row in preview_rows: preview_path = paths["state_dir"] / preview_row["rel_preview_path"] if not preview_path.exists(): continue current_html = preview_path.read_text(encoding="utf-8") links = build_document_attachment_preview_links(paths, connection, preview_path, child_rows) updated_html = inject_html_preview_attachment_links(current_html, links) if updated_html == current_html: continue preview_path.write_text(updated_html, encoding="utf-8") updated += 1 return updated def regenerate_production_preview_for_document( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, text_content: str, ) -> dict[str, object]: """Rebuild the synthesized HTML preview for a production document. Returns a status payload. Never raises on expected skips (non-production doc, no HTML preview row, missing production row). Unexpected I/O errors propagate and should be handled by the caller as a best-effort regen. """ document_row = connection.execute( """ SELECT id, control_number, production_id, begin_bates, end_bates, begin_attachment, end_attachment, title FROM documents WHERE id = ? """, (document_id,), ).fetchone() if document_row is None: return {"status": "skipped", "reason": "unknown_document"} if document_row["production_id"] is None: return {"status": "skipped", "reason": "not_production"} preview_row = connection.execute( """ SELECT rel_preview_path FROM document_previews WHERE document_id = ? AND preview_type = 'html' ORDER BY ordinal LIMIT 1 """, (document_id,), ).fetchone() if preview_row is None: return {"status": "skipped", "reason": "no_html_preview"} production_row = connection.execute( "SELECT production_name FROM productions WHERE id = ?", (document_row["production_id"],), ).fetchone() if production_row is None: return {"status": "skipped", "reason": "missing_production_row"} image_rows = connection.execute( """ SELECT ordinal, label, rel_source_path FROM document_source_parts WHERE document_id = ? AND part_kind = 'image' ORDER BY ordinal """, (document_id,), ).fetchall() page_images: list[dict[str, object]] = [] for index, row in enumerate(image_rows, start=1): abs_image = paths["root"] / row["rel_source_path"] if not abs_image.exists(): continue data_url = image_path_data_url( abs_image, max_dimension=INGEST_V2_PRODUCTION_PREVIEW_IMAGE_MAX_DIMENSION, ) if data_url is None: continue page_images.append( { "label": row["label"] or f"Page {index}", "src": data_url, } ) resolved_title = document_row["title"] or document_row["control_number"] html_body = build_production_preview_html( document_title=resolved_title, control_number=document_row["control_number"], production_name=production_row["production_name"], begin_bates=document_row["begin_bates"] or "", end_bates=document_row["end_bates"] or "", begin_attachment=document_row["begin_attachment"], end_attachment=document_row["end_attachment"], text_content=text_content, page_images=page_images, ) preview_path = paths["state_dir"] / preview_row["rel_preview_path"] preview_path.parent.mkdir(parents=True, exist_ok=True) preview_path.write_text(html_body, encoding="utf-8") sync_document_attachment_preview_links(connection, paths, document_id) return { "status": "ok", "rel_preview_path": preview_row["rel_preview_path"], "page_images": len(page_images), "text_chars": len(text_content), } def replace_document_related_rows( connection: sqlite3.Connection, document_id: int, metadata_values: dict[str, object], chunks: list[dict[str, object]], preview_rows: list[dict[str, object]], ) -> None: delete_document_related_rows(connection, document_id) row = connection.execute( "SELECT custodians_json FROM documents WHERE id = ?", (document_id,), ).fetchone() custodian_text = document_custodian_display_text_from_row(row) insert_document_preview_rows(connection, document_id, preview_rows) replace_document_chunks(connection, document_id, chunks) connection.execute( """ INSERT INTO documents_fts (document_id, file_name, title, subject, author, custodian, participants, recipients) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( document_id, metadata_values["file_name"], metadata_values["title"], metadata_values["subject"], metadata_values["author"], custodian_text, metadata_values["participants"], metadata_values["recipients"], ), ) def insert_document_preview_rows( connection: sqlite3.Connection, document_id: int, preview_rows: list[dict[str, object]], ) -> None: if not preview_rows: return connection.executemany( """ INSERT INTO document_previews ( document_id, rel_preview_path, preview_type, target_fragment, label, ordinal, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, [ ( document_id, row["rel_preview_path"], row["preview_type"], row.get("target_fragment"), row.get("label"), row["ordinal"], row["created_at"], ) for row in preview_rows ], ) def replace_document_preview_rows( connection: sqlite3.Connection, document_id: int, preview_rows: list[dict[str, object]], ) -> None: connection.execute("DELETE FROM document_previews WHERE document_id = ?", (document_id,)) insert_document_preview_rows(connection, document_id, preview_rows) def replace_document_source_parts( connection: sqlite3.Connection, document_id: int, source_parts: list[dict[str, object]], ) -> None: connection.execute("DELETE FROM document_source_parts WHERE document_id = ?", (document_id,)) if not source_parts: return connection.executemany( """ INSERT INTO document_source_parts ( document_id, part_kind, rel_source_path, ordinal, label, created_at ) VALUES (?, ?, ?, ?, ?, ?) """, [ ( document_id, row["part_kind"], row["rel_source_path"], row.get("ordinal", 0), row.get("label"), row.get("created_at", utc_now()), ) for row in source_parts ], ) def replace_document_chunks( connection: sqlite3.Connection, document_id: int, chunks: list[dict[str, object]], ) -> None: connection.execute("DELETE FROM document_chunks WHERE document_id = ?", (document_id,)) connection.execute("DELETE FROM chunks_fts WHERE document_id = ?", (document_id,)) if not chunks: return connection.executemany( """ INSERT INTO document_chunks ( document_id, chunk_index, char_start, char_end, token_estimate, text_content ) VALUES (?, ?, ?, ?, ?, ?) """, [ ( document_id, chunk["chunk_index"], chunk["char_start"], chunk["char_end"], chunk["token_estimate"], chunk["text_content"], ) for chunk in chunks ], ) chunk_rows = connection.execute( """ SELECT id, document_id, text_content FROM document_chunks WHERE document_id = ? ORDER BY chunk_index ASC """, (document_id,), ).fetchall() connection.executemany( """ INSERT INTO chunks_fts (chunk_id, document_id, text_content) VALUES (?, ?, ?) """, [(row["id"], row["document_id"], row["text_content"]) for row in chunk_rows], ) def delete_document_related_rows(connection: sqlite3.Connection, document_id: int) -> None: connection.execute("DELETE FROM document_chunks WHERE document_id = ?", (document_id,)) connection.execute("DELETE FROM chunks_fts WHERE document_id = ?", (document_id,)) connection.execute("DELETE FROM document_previews WHERE document_id = ?", (document_id,)) connection.execute("DELETE FROM document_source_parts WHERE document_id = ?", (document_id,)) connection.execute("DELETE FROM documents_fts WHERE document_id = ?", (document_id,)) def cleanup_unreferenced_preview_files( paths: dict[str, Path], connection: sqlite3.Connection, rel_preview_paths: list[str] | set[str] | tuple[str, ...], ) -> None: seen: set[str] = set() for rel_preview_path in rel_preview_paths: normalized_rel_path = normalize_whitespace(str(rel_preview_path or "")) if not normalized_rel_path or normalized_rel_path in seen: continue seen.add(normalized_rel_path) referenced_elsewhere = connection.execute( """ SELECT 1 FROM document_previews WHERE rel_preview_path = ? LIMIT 1 """, (normalized_rel_path,), ).fetchone() if referenced_elsewhere is None: remove_file_if_exists(paths["state_dir"] / normalized_rel_path) def prune_empty_conversation_preview_dirs(paths: dict[str, Path]) -> int: conversations_dir = paths["state_dir"] / "previews" / "conversations" if not conversations_dir.exists(): return 0 pruned = 0 for preview_dir in sorted(conversations_dir.glob("conversation-*"), reverse=True): if not preview_dir.is_dir(): continue try: preview_dir.rmdir() except OSError: continue pruned += 1 return pruned def refresh_documents_fts_row(connection: sqlite3.Connection, document_id: int) -> None: connection.execute("DELETE FROM documents_fts WHERE document_id = ?", (document_id,)) row = connection.execute( """ SELECT id, file_name, title, subject, author, custodians_json, participants, recipients FROM documents WHERE id = ? """, (document_id,), ).fetchone() if row is None: return connection.execute( """ INSERT INTO documents_fts (document_id, file_name, title, subject, author, custodian, participants, recipients) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( row["id"], row["file_name"], row["title"], row["subject"], row["author"], document_custodian_display_text_from_row(row), row["participants"], row["recipients"], ), ) PREVIEW_ARTIFACT_WRITE_RETRY_DELAYS_SECONDS = (0.05, 0.15, 0.35) def preview_artifact_write_is_retryable(exc: OSError) -> bool: return isinstance(exc, PermissionError) or getattr(exc, "errno", None) in { errno.EACCES, errno.EAGAIN, errno.EBUSY, errno.EPERM, } def write_preview_artifact_text_atomic(path: Path, content: str) -> None: retry_delays = PREVIEW_ARTIFACT_WRITE_RETRY_DELAYS_SECONDS for attempt in range(len(retry_delays) + 1): temp_path: Path | None = None try: path.parent.mkdir(parents=True, exist_ok=True) temp_file = tempfile.NamedTemporaryFile( mode="w", encoding="utf-8", dir=str(path.parent), prefix=f".{path.name}.", suffix=".tmp", delete=False, ) temp_path = Path(temp_file.name) with temp_file: temp_file.write(content) temp_file.flush() os.replace(temp_path, path) return except OSError as exc: if temp_path is not None: try: temp_path.unlink() except FileNotFoundError: pass except OSError: pass if attempt >= len(retry_delays) or not preview_artifact_write_is_retryable(exc): raise time.sleep(retry_delays[attempt]) def write_preview_artifacts( paths: dict[str, Path], rel_path: str, preview_artifacts: list[dict[str, object]] ) -> list[dict[str, object]]: preview_base = preview_base_path_for_rel_path(rel_path) preview_rows: list[dict[str, object]] = [] for artifact in preview_artifacts: preview_rel_path = preview_base / str(artifact["file_name"]) absolute_path = paths["state_dir"] / preview_rel_path write_preview_artifact_text_atomic(absolute_path, str(artifact["content"])) preview_rows.append( { "rel_preview_path": preview_rel_path.as_posix(), "preview_type": artifact["preview_type"], "target_fragment": artifact.get("target_fragment"), "label": artifact.get("label"), "ordinal": int(artifact.get("ordinal", 0)), "created_at": utc_now(), } ) return preview_rows def write_preview_page_assets( paths: dict[str, Path], page_assets: list[dict[str, object]], ) -> list[dict[str, object]]: preview_rows: list[dict[str, object]] = [] for asset in page_assets: rel_preview_path = normalize_whitespace(str(asset.get("rel_preview_path") or "")) payload = asset.get("payload") if not rel_preview_path or not isinstance(payload, (bytes, bytearray)): continue absolute_path = paths["state_dir"] / rel_preview_path absolute_path.parent.mkdir(parents=True, exist_ok=True) absolute_path.write_bytes(bytes(payload)) preview_rows.append( { "rel_preview_path": rel_preview_path, "preview_type": "image", "target_fragment": None, "label": asset.get("label"), "ordinal": int(asset.get("ordinal") or 0), "created_at": utc_now(), } ) return preview_rows def write_attachment_blob( paths: dict[str, Path], parent_rel_path: str, control_number: str, file_name: str, payload: bytes, ) -> tuple[str, Path]: storage_name = sanitize_storage_filename(file_name) source_rel_path = container_source_rel_path_from_message_rel_path(parent_rel_path) if source_rel_path is not None: preview_rel_path = Path("previews") / Path(source_rel_path) / "attachments" / control_number / storage_name else: preview_rel_path = Path("previews") / Path(parent_rel_path) / "attachments" / control_number / storage_name absolute_path = paths["state_dir"] / preview_rel_path absolute_path.parent.mkdir(parents=True, exist_ok=True) absolute_path.write_bytes(payload) rel_path = Path(INTERNAL_REL_PATH_PREFIX) / preview_rel_path return rel_path.as_posix(), absolute_path def remove_file_if_exists(path: Path, *, ignore_permission_error: bool = False) -> bool: retry_delays = PREVIEW_ARTIFACT_WRITE_RETRY_DELAYS_SECONDS for attempt in range(len(retry_delays) + 1): try: if path.is_file() or path.is_symlink(): path.unlink() return True return False except FileNotFoundError: return False except OSError as exc: if attempt < len(retry_delays) and preview_artifact_write_is_retryable(exc): time.sleep(retry_delays[attempt]) continue if ignore_permission_error and preview_artifact_write_is_retryable(exc): return False raise def cleanup_document_artifacts( paths: dict[str, Path], connection: sqlite3.Connection, row: sqlite3.Row | None, ) -> None: if row is None: return preview_rows = connection.execute( """ SELECT rel_preview_path FROM document_previews WHERE document_id = ? """, (row["id"],), ).fetchall() for preview_row in preview_rows: referenced_elsewhere = connection.execute( """ SELECT 1 FROM document_previews WHERE rel_preview_path = ? AND document_id != ? LIMIT 1 """, (preview_row["rel_preview_path"], row["id"]), ).fetchone() if referenced_elsewhere is None: remove_file_if_exists( paths["state_dir"] / preview_row["rel_preview_path"], ignore_permission_error=True, ) if is_internal_rel_path(row["rel_path"]): remove_file_if_exists(document_absolute_path(paths, row["rel_path"])) def parse_file_types(raw_value: str | None) -> set[str] | None: if raw_value is None: return None parts = {part.strip().lower() for part in raw_value.split(",") if part.strip()} invalid = sorted(parts - SUPPORTED_FILE_TYPES) if invalid: raise RetrieverError(f"Unsupported file type filters: {', '.join(invalid)}") return parts def apply_manual_locks(existing_row: sqlite3.Row | None, extracted: dict[str, object]) -> dict[str, object]: if existing_row is None: return extracted locked_fields = set(normalize_string_list(existing_row[MANUAL_FIELD_LOCKS_COLUMN])) merged = dict(extracted) for field_name in locked_fields: if field_name in EDITABLE_BUILTIN_FIELDS: merged[field_name] = existing_row[field_name] return merged def build_fallback_extract(path: Path) -> dict[str, object]: return { "page_count": None, "author": None, "content_type": infer_content_type_from_extension(normalize_extension(path)), "custodian": None, "date_created": None, "date_modified": None, "participants": None, "title": path.stem or path.name, "subject": None, "recipients": None, "text_content": "", "text_status": "empty", "preview_artifacts": [], "attachments": [], } def upsert_document_row( connection: sqlite3.Connection, rel_path: str, source_path: Path | None, existing_row: sqlite3.Row | None, extracted: dict[str, object], *, existing_occurrence_row: sqlite3.Row | None = None, file_name: str, parent_document_id: int | None, child_document_kind: str | None = None, control_number: str, dataset_id: int | None, conversation_id: int | None = None, conversation_assignment_mode: str | None = None, control_number_batch: int | None, control_number_family_sequence: int | None, control_number_attachment_sequence: int | None, root_message_key: str | None = None, source_kind: str | None = None, source_rel_path: str | None = None, source_item_id: str | None = None, source_folder_path: str | None = None, custodian_override: str | None = None, production_id: int | None = None, begin_bates: str | None = None, end_bates: str | None = None, begin_attachment: str | None = None, end_attachment: str | None = None, file_type_override: str | None = None, file_size_override: int | None = None, file_hash_override: str | None = None, ingested_at_override: str | None = None, last_seen_at_override: str | None = None, updated_at_override: str | None = None, ) -> int: now = utc_now() content_hash = sha256_text(str(extracted["text_content"] or "")) file_hash = file_hash_override if file_hash_override is not None else (sha256_file(source_path) if source_path is not None else None) file_size = file_size_override if file_size_override is not None else (source_path.stat().st_size if source_path is not None and source_path.exists() else None) file_type = file_type_override or normalize_extension(Path(file_name)) or (normalize_extension(source_path) if source_path is not None else None) effective_child_kind = effective_child_document_kind( parent_document_id=parent_document_id, child_document_kind=( child_document_kind if child_document_kind is not None else existing_row["child_document_kind"] if existing_row is not None and "child_document_kind" in existing_row.keys() else None ), ) effective_source_kind = source_kind or ( EMAIL_ATTACHMENT_SOURCE_KIND if effective_child_kind == CHILD_DOCUMENT_KIND_ATTACHMENT else FILESYSTEM_SOURCE_KIND ) effective_source_rel_path = source_rel_path if effective_source_rel_path is None: if existing_row is not None and existing_row["source_rel_path"] is not None: effective_source_rel_path = existing_row["source_rel_path"] else: effective_source_rel_path = rel_path effective_dataset_id = dataset_id if effective_dataset_id is None and existing_row is not None and existing_row["dataset_id"] is not None: effective_dataset_id = int(existing_row["dataset_id"]) normalized_extracted_content_type = normalize_whitespace(str(extracted.get("content_type") or "")).lower() should_preserve_existing_conversation = ( effective_child_kind is not None or normalized_extracted_content_type in {"email", "chat"} or conversation_id is not None ) effective_conversation_id = conversation_id if ( effective_conversation_id is None and should_preserve_existing_conversation and existing_row is not None and "conversation_id" in existing_row.keys() ): if existing_row["conversation_id"] is not None: effective_conversation_id = int(existing_row["conversation_id"]) effective_conversation_assignment = effective_conversation_assignment_mode( conversation_assignment_mode if conversation_assignment_mode is not None else existing_row["conversation_assignment_mode"] if existing_row is not None and "conversation_assignment_mode" in existing_row.keys() else None ) if not should_preserve_existing_conversation: effective_conversation_assignment = CONVERSATION_ASSIGNMENT_MODE_AUTO effective_root_message_key = root_message_key if effective_root_message_key is None and existing_row is not None and "root_message_key" in existing_row.keys(): effective_root_message_key = existing_row["root_message_key"] custodian = extracted.get("custodian") if custodian is None: custodian = infer_source_custodian( source_kind=effective_source_kind, source_rel_path=source_rel_path, parent_custodian=custodian_override, ) common_values = { "control_number": control_number, "canonical_kind": canonical_kind_from_metadata( extracted_content_type=extracted.get("content_type"), file_type=file_type, source_kind=effective_source_kind, ), "canonical_status": CANONICAL_STATUS_ACTIVE, "conversation_id": effective_conversation_id, "conversation_assignment_mode": effective_conversation_assignment, "dataset_id": effective_dataset_id, "parent_document_id": parent_document_id, "child_document_kind": effective_child_kind, "source_kind": effective_source_kind, "source_rel_path": effective_source_rel_path, "source_item_id": source_item_id, "root_message_key": effective_root_message_key, "source_folder_path": source_folder_path, "production_id": production_id, "begin_bates": begin_bates, "end_bates": end_bates, "begin_attachment": begin_attachment, "end_attachment": end_attachment, "rel_path": rel_path, "file_name": file_name, "file_type": file_type, "file_size": file_size, "page_count": extracted.get("page_count"), "author": extracted.get("author"), "content_type": extracted.get("content_type"), "date_created": extracted.get("date_created"), "date_modified": extracted.get("date_modified"), "participants": extracted.get("participants"), "title": extracted.get("title"), "subject": extracted.get("subject"), "recipients": extracted.get("recipients"), "file_hash": file_hash, "content_hash": content_hash, "text_status": extracted.get("text_status", "ok"), "lifecycle_status": "active", "ingested_at": ingested_at_override or now, "last_seen_at": last_seen_at_override or now, "updated_at": updated_at_override or now, "control_number_batch": control_number_batch, "control_number_family_sequence": control_number_family_sequence, "control_number_attachment_sequence": control_number_attachment_sequence, } if existing_row is None: connection.execute( """ INSERT INTO documents ( control_number, canonical_kind, canonical_status, merged_into_document_id, conversation_id, conversation_assignment_mode, dataset_id, parent_document_id, child_document_kind, source_kind, source_rel_path, source_item_id, root_message_key, source_folder_path, production_id, begin_bates, end_bates, begin_attachment, end_attachment, rel_path, file_name, file_type, file_size, page_count, author, date_created, content_type, date_modified, title, subject, participants, recipients, manual_field_locks_json, file_hash, content_hash, text_status, lifecycle_status, ingested_at, last_seen_at, updated_at, control_number_batch, control_number_family_sequence, control_number_attachment_sequence ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( common_values["control_number"], common_values["canonical_kind"], common_values["canonical_status"], None, common_values["conversation_id"], common_values["conversation_assignment_mode"], common_values["dataset_id"], common_values["parent_document_id"], common_values["child_document_kind"], common_values["source_kind"], common_values["source_rel_path"], common_values["source_item_id"], common_values["root_message_key"], common_values["source_folder_path"], common_values["production_id"], common_values["begin_bates"], common_values["end_bates"], common_values["begin_attachment"], common_values["end_attachment"], common_values["rel_path"], common_values["file_name"], common_values["file_type"], common_values["file_size"], common_values["page_count"], common_values["author"], common_values["date_created"], common_values["content_type"], common_values["date_modified"], common_values["title"], common_values["subject"], common_values["participants"], common_values["recipients"], "[]", common_values["file_hash"], common_values["content_hash"], common_values["text_status"], common_values["lifecycle_status"], common_values["ingested_at"], common_values["last_seen_at"], common_values["updated_at"], common_values["control_number_batch"], common_values["control_number_family_sequence"], common_values["control_number_attachment_sequence"], ), ) document_id = int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) else: locked_value = existing_row[MANUAL_FIELD_LOCKS_COLUMN] or "[]" connection.execute( """ UPDATE documents SET control_number = ?, canonical_kind = ?, canonical_status = ?, merged_into_document_id = NULL, conversation_id = ?, conversation_assignment_mode = ?, dataset_id = ?, parent_document_id = ?, child_document_kind = ?, source_kind = ?, source_rel_path = ?, source_item_id = ?, root_message_key = ?, source_folder_path = ?, production_id = ?, begin_bates = ?, end_bates = ?, begin_attachment = ?, end_attachment = ?, rel_path = ?, file_name = ?, file_type = ?, file_size = ?, page_count = ?, author = ?, content_type = ?, date_created = ?, date_modified = ?, title = ?, subject = ?, participants = ?, recipients = ?, file_hash = ?, content_hash = ?, text_status = ?, lifecycle_status = ?, ingested_at = ?, last_seen_at = ?, updated_at = ?, manual_field_locks_json = ?, control_number_batch = ?, control_number_family_sequence = ?, control_number_attachment_sequence = ? WHERE id = ? """, ( common_values["control_number"], common_values["canonical_kind"], common_values["canonical_status"], common_values["conversation_id"], common_values["conversation_assignment_mode"], common_values["dataset_id"], common_values["parent_document_id"], common_values["child_document_kind"], common_values["source_kind"], common_values["source_rel_path"], common_values["source_item_id"], common_values["root_message_key"], common_values["source_folder_path"], common_values["production_id"], common_values["begin_bates"], common_values["end_bates"], common_values["begin_attachment"], common_values["end_attachment"], common_values["rel_path"], common_values["file_name"], common_values["file_type"], common_values["file_size"], common_values["page_count"], common_values["author"], common_values["content_type"], common_values["date_created"], common_values["date_modified"], common_values["title"], common_values["subject"], common_values["participants"], common_values["recipients"], common_values["file_hash"], common_values["content_hash"], common_values["text_status"], common_values["lifecycle_status"], common_values["ingested_at"], common_values["last_seen_at"], common_values["updated_at"], locked_value, common_values["control_number_batch"], common_values["control_number_family_sequence"], common_values["control_number_attachment_sequence"], existing_row["id"], ), ) document_id = int(existing_row["id"]) effective_existing_occurrence = existing_occurrence_row if effective_existing_occurrence is None and existing_row is not None: effective_existing_occurrence = find_active_occurrence_by_source_identity( connection, source_kind=effective_source_kind, custodian=custodian, source_rel_path=effective_source_rel_path, source_item_id=source_item_id, ) if effective_existing_occurrence is None: effective_existing_occurrence = connection.execute( """ SELECT * FROM document_occurrences WHERE document_id = ? AND rel_path = ? AND lifecycle_status = ? ORDER BY id ASC LIMIT 1 """, (document_id, rel_path, ACTIVE_OCCURRENCE_STATUS), ).fetchone() parent_occurrence_id = None if parent_document_id is not None: parent_occurrence = select_preferred_occurrence(active_occurrence_rows_for_document(connection, int(parent_document_id))) if parent_occurrence is not None: parent_occurrence_id = int(parent_occurrence["id"]) upsert_document_occurrence( connection, document_id=document_id, existing_occurrence_id=(int(effective_existing_occurrence["id"]) if effective_existing_occurrence is not None else None), parent_occurrence_id=parent_occurrence_id, occurrence_control_number=control_number, source_kind=effective_source_kind, source_rel_path=effective_source_rel_path, source_item_id=source_item_id, source_folder_path=source_folder_path, production_id=production_id, begin_bates=begin_bates, end_bates=end_bates, begin_attachment=begin_attachment, end_attachment=end_attachment, rel_path=rel_path, file_name=file_name, file_type=file_type, mime_type=None, file_size=file_size, file_hash=file_hash, custodian=custodian, fs_created_at=None, fs_modified_at=None, extracted=extracted, has_preview=bool(extracted.get("preview_artifacts")), text_status=str(common_values["text_status"]), ingested_at=str(common_values["ingested_at"]), last_seen_at=str(common_values["last_seen_at"]), updated_at=str(common_values["updated_at"]), ) if parent_document_id is None and effective_source_kind in {FILESYSTEM_SOURCE_KIND, PST_SOURCE_KIND, MBOX_SOURCE_KIND}: bind_document_dedupe_key( connection, basis="file_hash", key_value=file_hash, document_id=document_id, ) refresh_source_backed_dataset_memberships_for_document(connection, document_id) refresh_document_from_occurrences(connection, document_id) return document_id def attach_occurrence_to_existing_document( connection: sqlite3.Connection, document_row: sqlite3.Row, *, existing_occurrence_row: sqlite3.Row | None, rel_path: str, file_name: str, file_type: str | None, file_size: int | None, file_hash: str | None, source_kind: str, source_rel_path: str, source_item_id: str | None, source_folder_path: str | None, custodian: str | None, production_id: int | None = None, begin_bates: str | None = None, end_bates: str | None = None, begin_attachment: str | None = None, end_attachment: str | None = None, parent_document_id: int | None = None, parent_occurrence_id: int | None = None, occurrence_control_number: str | None = None, ingested_at: str | None = None, last_seen_at: str | None = None, updated_at: str | None = None, ) -> int: resolved_parent_occurrence_id = parent_occurrence_id if resolved_parent_occurrence_id is None and parent_document_id is not None: parent_occurrence = select_preferred_occurrence(active_occurrence_rows_for_document(connection, parent_document_id)) if parent_occurrence is not None: resolved_parent_occurrence_id = int(parent_occurrence["id"]) preview_exists = connection.execute( "SELECT 1 FROM document_previews WHERE document_id = ? LIMIT 1", (document_row["id"],), ).fetchone() is not None occurrence_id = upsert_document_occurrence( connection, document_id=int(document_row["id"]), existing_occurrence_id=(int(existing_occurrence_row["id"]) if existing_occurrence_row is not None else None), parent_occurrence_id=resolved_parent_occurrence_id, occurrence_control_number=occurrence_control_number or document_row["control_number"], source_kind=source_kind, source_rel_path=source_rel_path, source_item_id=source_item_id, source_folder_path=source_folder_path, production_id=production_id, begin_bates=begin_bates, end_bates=end_bates, begin_attachment=begin_attachment, end_attachment=end_attachment, rel_path=rel_path, file_name=file_name, file_type=file_type, mime_type=None, file_size=file_size, file_hash=file_hash, custodian=custodian, fs_created_at=None, fs_modified_at=None, extracted={ "author": document_row["author"], "title": document_row["title"], "subject": document_row["subject"], "participants": document_row["participants"], "recipients": document_row["recipients"], "date_created": document_row["date_created"], "date_modified": document_row["date_modified"], "content_type": document_row["content_type"], }, has_preview=preview_exists, text_status=str(document_row["text_status"] or "ok"), ingested_at=ingested_at or str(document_row["ingested_at"] or utc_now()), last_seen_at=last_seen_at or str(document_row["last_seen_at"] or document_row["ingested_at"] or utc_now()), updated_at=updated_at or str(document_row["updated_at"] or utc_now()), ) if parent_document_id is None and source_kind in {FILESYSTEM_SOURCE_KIND, PST_SOURCE_KIND, MBOX_SOURCE_KIND}: bind_document_dedupe_key( connection, basis="file_hash", key_value=file_hash, document_id=int(document_row["id"]), ) refresh_source_backed_dataset_memberships_for_document(connection, int(document_row["id"])) refresh_document_from_occurrences(connection, int(document_row["id"])) return occurrence_id def clone_duplicate_family_child_occurrences( connection: sqlite3.Connection, paths: dict[str, Path], *, parent_document_id: int, parent_occurrence_id: int, parent_rel_path: str, custodian: str | None, ingested_at: str | None = None, last_seen_at: str | None = None, updated_at: str | None = None, ) -> int: child_rows = connection.execute( """ SELECT * FROM documents WHERE parent_document_id = ? AND lifecycle_status != 'deleted' ORDER BY CASE WHEN control_number_attachment_sequence IS NULL THEN 1 ELSE 0 END ASC, control_number_attachment_sequence ASC, id ASC """, (parent_document_id,), ).fetchall() cloned_count = 0 for child_row in child_rows: child_document_id = int(child_row["id"]) preferred_occurrence = select_preferred_occurrence( active_occurrence_rows_for_document(connection, child_document_id) ) child_source_rel_path = ( str(preferred_occurrence["rel_path"]) if preferred_occurrence is not None and preferred_occurrence["rel_path"] is not None else str(child_row["rel_path"]) ) child_source_path = document_absolute_path(paths, child_source_rel_path) if not child_source_path.exists(): raise RetrieverError( f"Could not clone duplicate family child {child_document_id}: missing source artifact {child_source_rel_path!r}." ) cloned_rel_path, cloned_path = write_attachment_blob( paths, parent_rel_path, str(child_row["control_number"]), str(child_row["file_name"]), child_source_path.read_bytes(), ) attach_occurrence_to_existing_document( connection, child_row, existing_occurrence_row=None, rel_path=cloned_rel_path, file_name=str(child_row["file_name"]), file_type=( str(preferred_occurrence["file_type"]) if preferred_occurrence is not None and preferred_occurrence["file_type"] is not None else child_row["file_type"] ), file_size=( int(preferred_occurrence["file_size"]) if preferred_occurrence is not None and preferred_occurrence["file_size"] is not None else cloned_path.stat().st_size ), file_hash=( str(preferred_occurrence["file_hash"]) if preferred_occurrence is not None and preferred_occurrence["file_hash"] is not None else str(child_row["file_hash"]) if child_row["file_hash"] is not None else sha256_file(cloned_path) ), source_kind=( str(preferred_occurrence["source_kind"]) if preferred_occurrence is not None and preferred_occurrence["source_kind"] is not None else str(child_row["source_kind"]) ), source_rel_path=cloned_rel_path, source_item_id=( str(preferred_occurrence["source_item_id"]) if preferred_occurrence is not None and preferred_occurrence["source_item_id"] is not None else str(child_row["source_item_id"]) if child_row["source_item_id"] is not None else None ), source_folder_path=( str(preferred_occurrence["source_folder_path"]) if preferred_occurrence is not None and preferred_occurrence["source_folder_path"] is not None else str(child_row["source_folder_path"]) if child_row["source_folder_path"] is not None else None ), custodian=custodian, parent_document_id=parent_document_id, parent_occurrence_id=parent_occurrence_id, occurrence_control_number=str(child_row["control_number"] or ""), ingested_at=ingested_at, last_seen_at=last_seen_at, updated_at=updated_at, ) cloned_count += 1 return cloned_count def mark_seen_without_reingest( connection: sqlite3.Connection, occurrence_row: sqlite3.Row, *, dataset_id: int | None = None, dataset_source_id: int | None = None, ) -> None: now = utc_now() connection.execute( """ UPDATE document_occurrences SET lifecycle_status = 'active', last_seen_at = ?, updated_at = ? WHERE id = ? """, (now, now, occurrence_row["id"]), ) connection.execute( """ UPDATE document_occurrences SET lifecycle_status = 'active', last_seen_at = ?, updated_at = ? WHERE parent_occurrence_id = ? AND lifecycle_status != 'deleted' """, (now, now, occurrence_row["id"]), ) if dataset_id is not None: ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=int(occurrence_row["document_id"]), dataset_source_id=dataset_source_id, ) child_rows = connection.execute( """ SELECT DISTINCT document_id FROM document_occurrences WHERE parent_occurrence_id = ? AND lifecycle_status != 'deleted' ORDER BY document_id ASC """, (occurrence_row["id"],), ).fetchall() for child_row in child_rows: ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=int(child_row["document_id"]), dataset_source_id=dataset_source_id, ) refresh_source_backed_dataset_memberships_for_document(connection, int(child_row["document_id"])) refresh_document_from_occurrences(connection, int(child_row["document_id"])) refresh_source_backed_dataset_memberships_for_document(connection, int(occurrence_row["document_id"])) refresh_document_from_occurrences(connection, int(occurrence_row["document_id"])) def replace_document_email_threading_row( connection: sqlite3.Connection, *, document_id: int, email_threading: object, ) -> None: if not isinstance(email_threading, dict): connection.execute("DELETE FROM document_email_threading WHERE document_id = ?", (document_id,)) return references = [ value for value in normalize_string_list(email_threading.get("references")) if normalize_whitespace(value) ] normalized_payload = { "message_id": normalize_email_message_id(email_threading.get("message_id")), "in_reply_to": normalize_email_message_id(email_threading.get("in_reply_to")), "references_json": json.dumps(references), "conversation_index": normalize_whitespace(str(email_threading.get("conversation_index") or "")) or None, "conversation_topic": normalize_email_thread_subject(email_threading.get("conversation_topic")), "normalized_subject": normalize_email_thread_subject(email_threading.get("normalized_subject")), "updated_at": utc_now(), } if not any( ( normalized_payload["message_id"], normalized_payload["in_reply_to"], references, normalized_payload["conversation_index"], normalized_payload["conversation_topic"], normalized_payload["normalized_subject"], ) ): connection.execute("DELETE FROM document_email_threading WHERE document_id = ?", (document_id,)) return connection.execute( """ INSERT INTO document_email_threading ( document_id, message_id, in_reply_to, references_json, conversation_index, conversation_topic, normalized_subject, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(document_id) DO UPDATE SET message_id = excluded.message_id, in_reply_to = excluded.in_reply_to, references_json = excluded.references_json, conversation_index = excluded.conversation_index, conversation_topic = excluded.conversation_topic, normalized_subject = excluded.normalized_subject, updated_at = excluded.updated_at """, ( document_id, normalized_payload["message_id"], normalized_payload["in_reply_to"], normalized_payload["references_json"], normalized_payload["conversation_index"], normalized_payload["conversation_topic"], normalized_payload["normalized_subject"], normalized_payload["updated_at"], ), ) def replace_document_chat_threading_row( connection: sqlite3.Connection, *, document_id: int, chat_threading: object, ) -> None: if not isinstance(chat_threading, dict): connection.execute("DELETE FROM document_chat_threading WHERE document_id = ?", (document_id,)) return participant_names = sorted_unique_display_names(normalize_string_list(chat_threading.get("participants"))) normalized_payload = { "thread_id": normalize_pst_chat_thread_id(chat_threading.get("thread_id")), "message_id": normalize_pst_identifier(chat_threading.get("message_id")), "parent_message_id": normalize_pst_identifier(chat_threading.get("parent_message_id")), "thread_type": normalize_whitespace(str(chat_threading.get("thread_type") or "")).lower() or None, "participants_json": json.dumps(participant_names), "updated_at": utc_now(), } if not any( ( normalized_payload["thread_id"], normalized_payload["message_id"], normalized_payload["parent_message_id"], normalized_payload["thread_type"], participant_names, ) ): connection.execute("DELETE FROM document_chat_threading WHERE document_id = ?", (document_id,)) return connection.execute( """ INSERT INTO document_chat_threading ( document_id, thread_id, message_id, parent_message_id, thread_type, participants_json, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(document_id) DO UPDATE SET thread_id = excluded.thread_id, message_id = excluded.message_id, parent_message_id = excluded.parent_message_id, thread_type = excluded.thread_type, participants_json = excluded.participants_json, updated_at = excluded.updated_at """, ( document_id, normalized_payload["thread_id"], normalized_payload["message_id"], normalized_payload["parent_message_id"], normalized_payload["thread_type"], normalized_payload["participants_json"], normalized_payload["updated_at"], ), ) def document_row_has_email_threading(connection: sqlite3.Connection, row: sqlite3.Row | None) -> bool: if row is None: return False file_type = normalize_whitespace(str(row["file_type"] or "")).lower() if "file_type" in row.keys() else "" if file_type not in {"eml", "msg"}: return True signal_row = connection.execute( """ SELECT 1 FROM document_email_threading WHERE document_id = ? LIMIT 1 """, (int(row["id"]),), ).fetchone() return signal_row is not None def container_email_documents_missing_threading( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, ) -> bool: root_document_ids = sorted( container_root_document_ids_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ) if not root_document_ids: return False placeholders = ", ".join("?" for _ in root_document_ids) row = connection.execute( f""" SELECT 1 FROM documents d LEFT JOIN document_email_threading det ON det.document_id = d.id LEFT JOIN document_chat_threading dct ON dct.document_id = d.id WHERE d.id IN ({placeholders}) AND d.lifecycle_status != 'deleted' AND ( (d.content_type = 'Email' AND det.document_id IS NULL) OR (d.content_type = 'Chat' AND dct.document_id IS NULL) ) LIMIT 1 """, root_document_ids, ).fetchone() return row is not None def list_email_conversation_documents(connection: sqlite3.Connection) -> list[dict[str, object]]: rows = connection.execute( """ SELECT d.id, d.control_number, d.conversation_id, d.conversation_assignment_mode, d.source_kind, d.source_rel_path, d.source_folder_path, d.file_type, d.author, d.recipients, d.subject, d.title, d.date_created, d.custodians_json, det.message_id, det.in_reply_to, det.references_json, det.conversation_index, det.conversation_topic, det.normalized_subject FROM documents d LEFT JOIN document_email_threading det ON det.document_id = d.id WHERE d.parent_document_id IS NULL AND d.content_type = 'Email' AND d.lifecycle_status NOT IN ('missing', 'deleted') ORDER BY CASE WHEN d.date_created IS NULL OR TRIM(d.date_created) = '' THEN 1 ELSE 0 END ASC, d.date_created ASC, d.id ASC """ ).fetchall() documents: list[dict[str, object]] = [] for row in rows: references = normalize_string_list(row["references_json"]) normalized_subject = normalize_email_thread_subject( row["normalized_subject"] or row["subject"] or row["conversation_topic"] ) custodians = parse_document_custodians_json(row["custodians_json"]) documents.append( { "id": int(row["id"]), "control_number": row["control_number"], "existing_conversation_id": int(row["conversation_id"]) if row["conversation_id"] is not None else None, "assignment_mode": effective_conversation_assignment_mode(row["conversation_assignment_mode"]), "source_kind": normalize_whitespace(str(row["source_kind"] or "")).lower(), "source_rel_path": normalize_whitespace(str(row["source_rel_path"] or "")) or None, "source_folder_path": normalize_whitespace(str(row["source_folder_path"] or "")) or None, "file_type": normalize_whitespace(str(row["file_type"] or "")).lower() or None, "author": normalize_whitespace(str(row["author"] or "")) or None, "recipients": normalize_whitespace(str(row["recipients"] or "")) or None, "subject": normalize_whitespace(str(row["subject"] or "")) or None, "title": normalize_whitespace(str(row["title"] or "")) or None, "date_created": normalize_datetime(row["date_created"]), "custodians": custodians, "message_id": normalize_email_message_id(row["message_id"]), "in_reply_to": normalize_email_message_id(row["in_reply_to"]), "references": references, "conversation_index_root": normalize_email_conversation_index_root(row["conversation_index"]), "conversation_topic": normalize_email_thread_subject(row["conversation_topic"]), "normalized_subject": normalized_subject, "participant_keys": email_participant_keys(row["author"], row["recipients"]), "heuristic_scope": email_heuristic_scope_key(row["source_kind"], row["source_rel_path"]), } ) return documents def create_email_conversation_cluster(*, manual_conversation_id: int | None = None) -> dict[str, object]: return { "documents": [], "manual_conversation_id": manual_conversation_id, "message_ids": set(), "conversation_index_roots": set(), "conversation_topics": set(), "normalized_subjects": set(), "participant_keys": set(), "custodians": set(), "heuristic_scopes": set(), "latest_date": None, } def add_document_to_email_cluster( cluster: dict[str, object], document: dict[str, object], *, cluster_id: int, message_id_index: dict[str, set[int]], conversation_index_root_index: dict[str, set[int]], conversation_topic_index: dict[str, set[int]], heuristic_subject_index: dict[tuple[str, str], set[int]], ) -> None: cluster["documents"].append(document) message_id = document.get("message_id") if message_id: cluster["message_ids"].add(message_id) message_id_index.setdefault(str(message_id), set()).add(cluster_id) conversation_index_root = document.get("conversation_index_root") if conversation_index_root: cluster["conversation_index_roots"].add(conversation_index_root) conversation_index_root_index.setdefault(str(conversation_index_root), set()).add(cluster_id) conversation_topic = document.get("conversation_topic") if conversation_topic: cluster["conversation_topics"].add(conversation_topic) conversation_topic_index.setdefault(str(conversation_topic), set()).add(cluster_id) normalized_subject = document.get("normalized_subject") heuristic_scope = document.get("heuristic_scope") if normalized_subject: cluster["normalized_subjects"].add(normalized_subject) if heuristic_scope: cluster["heuristic_scopes"].add(heuristic_scope) if normalized_subject and heuristic_scope: heuristic_subject_index.setdefault((str(heuristic_scope), str(normalized_subject)), set()).add(cluster_id) participant_keys = set(document.get("participant_keys") or []) cluster["participant_keys"].update(participant_keys) cluster["custodians"].update(normalize_custodian_values(set(document.get("custodians") or []))) document_date = normalize_datetime(document.get("date_created")) if document_date and (cluster["latest_date"] is None or str(cluster["latest_date"]) < document_date): cluster["latest_date"] = document_date def choose_unique_cluster(cluster_ids: set[int]) -> int | None: if len(cluster_ids) == 1: return next(iter(cluster_ids)) return None def choose_reference_cluster(document: dict[str, object], message_id_index: dict[str, set[int]]) -> int | None: references = list(document.get("references") or []) if not references: return None candidate_cluster_ids: set[int] = set() for reference_id in references: candidate_cluster_ids.update(message_id_index.get(str(reference_id), set())) if not candidate_cluster_ids: return None scores: dict[int, int] = {} for cluster_id in candidate_cluster_ids: suffix_length = 0 for reference_id in reversed(references): if cluster_id in message_id_index.get(str(reference_id), set()): suffix_length += 1 elif suffix_length > 0: break if suffix_length > 0: scores[cluster_id] = suffix_length if not scores: return None best_score = max(scores.values()) best_cluster_ids = [cluster_id for cluster_id, score in scores.items() if score == best_score] if len(best_cluster_ids) != 1: return None return best_cluster_ids[0] def choose_outlook_cluster( document: dict[str, object], conversation_index_root_index: dict[str, set[int]], conversation_topic_index: dict[str, set[int]], ) -> int | None: conversation_index_root = document.get("conversation_index_root") conversation_topic = document.get("conversation_topic") if not conversation_index_root: return None root_candidates = ( set(conversation_index_root_index.get(str(conversation_index_root), set())) if conversation_index_root else set() ) if not root_candidates: return None if conversation_topic: topic_candidates = set(conversation_topic_index.get(str(conversation_topic), set())) intersection = root_candidates & topic_candidates chosen = choose_unique_cluster(intersection) if chosen is not None: return chosen return choose_unique_cluster(root_candidates) def email_document_has_strong_threading_signals(document: dict[str, object]) -> bool: if document.get("message_id"): return True if document.get("in_reply_to"): return True if list(document.get("references") or []): return True if document.get("conversation_index_root"): return True return False def email_document_allows_heuristic_fallback(document: dict[str, object]) -> bool: source_kind = normalize_whitespace(str(document.get("source_kind") or "")).lower() return source_kind in {FILESYSTEM_SOURCE_KIND, PRODUCTION_SOURCE_KIND} def email_document_can_use_heuristic_fallback(document: dict[str, object]) -> bool: return ( email_document_allows_heuristic_fallback(document) and not email_document_has_strong_threading_signals(document) ) def choose_heuristic_cluster( document: dict[str, object], clusters: list[dict[str, object]], heuristic_subject_index: dict[tuple[str, str], set[int]], ) -> int | None: normalized_subject = document.get("normalized_subject") heuristic_scope = document.get("heuristic_scope") if not normalized_subject or not heuristic_scope: return None candidate_cluster_ids = heuristic_subject_index.get((str(heuristic_scope), str(normalized_subject)), set()) if not candidate_cluster_ids: return None participant_keys = set(document.get("participant_keys") or []) document_date = normalize_datetime(document.get("date_created")) document_custodians = set(normalize_custodian_values(set(document.get("custodians") or []))) scored_candidates: list[tuple[int, str, int]] = [] for cluster_id in candidate_cluster_ids: cluster = clusters[cluster_id] cluster_participants = set(cluster["participant_keys"]) overlap = len(cluster_participants & participant_keys) if overlap <= 0: continue cluster_custodians = set(cluster["custodians"]) if document_custodians and cluster_custodians and document_custodians.isdisjoint(cluster_custodians): continue latest_date = normalize_datetime(cluster["latest_date"]) if document_date and latest_date and latest_date > document_date: continue scored_candidates.append((overlap, latest_date or "", cluster_id)) if not scored_candidates: return None scored_candidates.sort(key=lambda item: (item[0], item[1], -item[2]), reverse=True) best_candidate = scored_candidates[0] if len(scored_candidates) > 1 and scored_candidates[1][:2] == best_candidate[:2]: return None return best_candidate[2] def derive_email_conversation_key(cluster: dict[str, object]) -> str: message_ids = sorted(str(value) for value in cluster["message_ids"] if normalize_whitespace(str(value))) if message_ids: return f"rfc:{message_ids[0]}" conversation_index_roots = sorted( str(value) for value in cluster["conversation_index_roots"] if normalize_whitespace(str(value)) ) conversation_topics = sorted( str(value) for value in cluster["conversation_topics"] if normalize_whitespace(str(value)) ) if conversation_index_roots and conversation_topics: return f"outlook:{conversation_topics[0]}:{conversation_index_roots[0]}" if conversation_index_roots: return f"outlook_index:{conversation_index_roots[0]}" normalized_subjects = sorted( str(value) for value in cluster["normalized_subjects"] if normalize_whitespace(str(value)) ) heuristic_scopes = sorted( str(value) for value in cluster["heuristic_scopes"] if normalize_whitespace(str(value)) ) heuristic_signature = sha256_json_value( { "scope": heuristic_scopes[0] if heuristic_scopes else filesystem_dataset_locator(), "subject": normalized_subjects[0] if normalized_subjects else "", "participants": sorted(str(value) for value in cluster["participant_keys"]), } ) return f"heuristic:{heuristic_signature[:24]}" def derive_email_conversation_display_name(cluster: dict[str, object]) -> str: for document in list(cluster["documents"]): subject = normalize_email_thread_subject(document.get("subject"), preserve_case=True) if subject: return subject title = normalize_email_thread_subject(document.get("title"), preserve_case=True) if title: return title for conversation_topic in sorted( str(value) for value in cluster["conversation_topics"] if normalize_whitespace(str(value)) ): display_name = normalize_email_thread_subject(conversation_topic, preserve_case=True) if display_name: return display_name return "Email conversation" def sync_child_document_conversations( connection: sqlite3.Connection, *, parent_ids: set[int] | None = None, parent_content_types: set[str] | None = None, parent_source_kinds: set[str] | None = None, ) -> int: clauses = [ "parent.lifecycle_status NOT IN ('missing', 'deleted')", "child.lifecycle_status NOT IN ('missing', 'deleted')", ] parameters: list[object] = [] normalized_parent_ids = {int(value) for value in (parent_ids or set())} normalized_parent_content_types = sorted( normalize_whitespace(str(value or "")) for value in (parent_content_types or set()) if normalize_whitespace(str(value or "")) ) normalized_parent_source_kinds = sorted( normalize_whitespace(str(value or "")).lower() for value in (parent_source_kinds or set()) if normalize_whitespace(str(value or "")) ) if normalized_parent_ids: placeholders = ", ".join("?" for _ in normalized_parent_ids) clauses.append(f"parent.id IN ({placeholders})") parameters.extend(sorted(normalized_parent_ids)) if normalized_parent_content_types: placeholders = ", ".join("?" for _ in normalized_parent_content_types) clauses.append(f"parent.content_type IN ({placeholders})") parameters.extend(normalized_parent_content_types) if normalized_parent_source_kinds: placeholders = ", ".join("?" for _ in normalized_parent_source_kinds) clauses.append(f"parent.source_kind IN ({placeholders})") parameters.extend(normalized_parent_source_kinds) child_rows = connection.execute( f""" SELECT child.id, child.conversation_id AS child_conversation_id, child.conversation_assignment_mode AS child_conversation_assignment_mode, parent.conversation_id AS parent_conversation_id, parent.conversation_assignment_mode AS parent_conversation_assignment_mode FROM documents child JOIN documents parent ON parent.id = child.parent_document_id WHERE {' AND '.join(clauses)} ORDER BY child.id ASC """, tuple(parameters), ).fetchall() updated = 0 now = utc_now() for row in child_rows: parent_conversation_id = int(row["parent_conversation_id"]) if row["parent_conversation_id"] is not None else None parent_mode = effective_conversation_assignment_mode(row["parent_conversation_assignment_mode"]) child_conversation_id = int(row["child_conversation_id"]) if row["child_conversation_id"] is not None else None child_mode = effective_conversation_assignment_mode(row["child_conversation_assignment_mode"]) if child_conversation_id == parent_conversation_id and child_mode == parent_mode: continue connection.execute( """ UPDATE documents SET conversation_id = ?, conversation_assignment_mode = ?, updated_at = ? WHERE id = ? """, (parent_conversation_id, parent_mode, now, int(row["id"])), ) updated += 1 return updated def assign_email_conversations(connection: sqlite3.Connection) -> dict[str, int]: documents = list_email_conversation_documents(connection) if not documents: return { "email_conversations": 0, "email_documents_reassigned": 0, "email_child_documents_updated": 0, } clusters: list[dict[str, object]] = [] message_id_index: dict[str, set[int]] = {} conversation_index_root_index: dict[str, set[int]] = {} conversation_topic_index: dict[str, set[int]] = {} heuristic_subject_index: dict[tuple[str, str], set[int]] = {} manual_cluster_ids_by_conversation_id: dict[int, int] = {} for document in documents: if document["assignment_mode"] != CONVERSATION_ASSIGNMENT_MODE_MANUAL: continue existing_conversation_id = document.get("existing_conversation_id") if existing_conversation_id is None: continue cluster_id = manual_cluster_ids_by_conversation_id.get(int(existing_conversation_id)) if cluster_id is None: cluster_id = len(clusters) manual_cluster_ids_by_conversation_id[int(existing_conversation_id)] = cluster_id clusters.append(create_email_conversation_cluster(manual_conversation_id=int(existing_conversation_id))) add_document_to_email_cluster( clusters[cluster_id], document, cluster_id=cluster_id, message_id_index=message_id_index, conversation_index_root_index=conversation_index_root_index, conversation_topic_index=conversation_topic_index, heuristic_subject_index=heuristic_subject_index, ) auto_documents = [document for document in documents if document["assignment_mode"] != CONVERSATION_ASSIGNMENT_MODE_MANUAL] for document in auto_documents: cluster_id = None if document.get("message_id"): cluster_id = choose_unique_cluster(message_id_index.get(str(document["message_id"]), set())) if cluster_id is None and document.get("in_reply_to"): cluster_id = choose_unique_cluster(message_id_index.get(str(document["in_reply_to"]), set())) if cluster_id is None: cluster_id = choose_reference_cluster(document, message_id_index) if cluster_id is None: cluster_id = choose_outlook_cluster(document, conversation_index_root_index, conversation_topic_index) if cluster_id is None and email_document_can_use_heuristic_fallback(document): cluster_id = choose_heuristic_cluster(document, clusters, heuristic_subject_index) if cluster_id is None: cluster_id = len(clusters) clusters.append(create_email_conversation_cluster()) add_document_to_email_cluster( clusters[cluster_id], document, cluster_id=cluster_id, message_id_index=message_id_index, conversation_index_root_index=conversation_index_root_index, conversation_topic_index=conversation_topic_index, heuristic_subject_index=heuristic_subject_index, ) conversation_id_by_document_id: dict[int, int] = {} unique_conversation_ids: set[int] = set() for cluster in clusters: manual_conversation_id = cluster["manual_conversation_id"] if manual_conversation_id is not None: conversation_id = int(manual_conversation_id) else: conversation_id = upsert_conversation_row( connection, source_kind=EMAIL_CONVERSATION_SOURCE_KIND, source_locator=filesystem_dataset_locator(), conversation_key=derive_email_conversation_key(cluster), conversation_type="email", display_name=derive_email_conversation_display_name(cluster), ) unique_conversation_ids.add(conversation_id) for document in list(cluster["documents"]): conversation_id_by_document_id[int(document["id"])] = conversation_id documents_reassigned = 0 now = utc_now() for document in auto_documents: target_conversation_id = conversation_id_by_document_id.get(int(document["id"])) if target_conversation_id is None: continue if ( document.get("existing_conversation_id") == target_conversation_id and document.get("assignment_mode") == CONVERSATION_ASSIGNMENT_MODE_AUTO ): continue connection.execute( """ UPDATE documents SET conversation_id = ?, conversation_assignment_mode = ?, updated_at = ? WHERE id = ? AND COALESCE(conversation_assignment_mode, ?) != ? """, ( target_conversation_id, CONVERSATION_ASSIGNMENT_MODE_AUTO, now, int(document["id"]), CONVERSATION_ASSIGNMENT_MODE_AUTO, CONVERSATION_ASSIGNMENT_MODE_MANUAL, ), ) documents_reassigned += 1 child_documents_updated = sync_child_document_conversations(connection, parent_content_types={"Email"}) return { "email_conversations": len(unique_conversation_ids), "email_documents_reassigned": documents_reassigned, "email_child_documents_updated": child_documents_updated, } def list_pst_chat_conversation_documents(connection: sqlite3.Connection) -> list[dict[str, object]]: rows = connection.execute( """ SELECT d.id, d.control_number, d.conversation_id, d.conversation_assignment_mode, d.manual_field_locks_json, d.source_kind, d.source_rel_path, d.source_item_id, d.source_folder_path, d.file_type, d.title, d.participants, d.date_created, d.custodians_json, dct.thread_id, dct.thread_type, dct.participants_json FROM documents d LEFT JOIN document_chat_threading dct ON dct.document_id = d.id WHERE d.parent_document_id IS NULL AND d.source_kind = ? AND d.content_type = 'Chat' AND d.lifecycle_status NOT IN ('missing', 'deleted') ORDER BY CASE WHEN d.date_created IS NULL OR TRIM(d.date_created) = '' THEN 1 ELSE 0 END ASC, d.date_created ASC, d.id ASC """, (PST_SOURCE_KIND,), ).fetchall() documents: list[dict[str, object]] = [] for row in rows: participant_names = sorted_unique_display_names( [ *normalize_string_list(row["participants_json"]), *[ normalize_whitespace(part) for part in str(row["participants"] or "").split(",") if normalize_whitespace(part) ], ] ) documents.append( { "id": int(row["id"]), "control_number": row["control_number"], "existing_conversation_id": int(row["conversation_id"]) if row["conversation_id"] is not None else None, "assignment_mode": effective_conversation_assignment_mode(row["conversation_assignment_mode"]), "locked_fields": set(normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN])), "source_kind": normalize_whitespace(str(row["source_kind"] or "")).lower(), "source_rel_path": normalize_whitespace(str(row["source_rel_path"] or "")) or None, "source_item_id": normalize_whitespace(str(row["source_item_id"] or "")) or None, "source_folder_path": normalize_whitespace(str(row["source_folder_path"] or "")) or None, "file_type": normalize_whitespace(str(row["file_type"] or "")).lower() or None, "title": normalize_whitespace(str(row["title"] or "")) or None, "participants": normalize_whitespace(str(row["participants"] or "")) or None, "participant_names": participant_names, "date_created": normalize_datetime(row["date_created"]), "custodians": parse_document_custodians_json(row["custodians_json"]), "thread_id": normalize_pst_chat_thread_id(row["thread_id"]), "thread_type": normalize_whitespace(str(row["thread_type"] or "")).lower() or None, } ) return documents def create_pst_chat_conversation_cluster(*, manual_conversation_id: int | None = None) -> dict[str, object]: return { "documents": [], "manual_conversation_id": manual_conversation_id, "source_rel_paths": set(), "source_folder_paths": set(), "thread_ids": set(), "thread_types": set(), "participant_names": set(), "titles": set(), "participants": set(), } def add_document_to_pst_chat_cluster(cluster: dict[str, object], document: dict[str, object]) -> None: cluster["documents"].append(document) source_rel_path = document.get("source_rel_path") if source_rel_path: cluster["source_rel_paths"].add(source_rel_path) source_folder_path = document.get("source_folder_path") if source_folder_path: cluster["source_folder_paths"].add(source_folder_path) thread_id = document.get("thread_id") if thread_id: cluster["thread_ids"].add(thread_id) thread_type = document.get("thread_type") if thread_type: cluster["thread_types"].add(thread_type) for participant_name in list(document.get("participant_names") or []): if normalize_whitespace(str(participant_name)): cluster["participant_names"].add(str(participant_name)) title = document.get("title") if title: cluster["titles"].add(title) participants = document.get("participants") if participants: cluster["participants"].add(participants) def pst_chat_cluster_scope_key(document: dict[str, object]) -> tuple[str, str]: source_rel_path = normalize_whitespace(str(document.get("source_rel_path") or "")) thread_id = normalize_pst_chat_thread_id(document.get("thread_id")) if thread_id: return (filesystem_dataset_locator(), f"thread:{thread_id.lower()}") source_folder_path = normalize_whitespace(str(document.get("source_folder_path") or "")) if source_folder_path: return (source_rel_path, f"folder:{source_folder_path.lower()}") source_item_id = normalize_whitespace(str(document.get("source_item_id") or "")) if source_item_id: return (source_rel_path, f"item:{source_item_id.lower()}") return (source_rel_path, f"doc:{int(document['id'])}") def derive_pst_chat_conversation_key(cluster: dict[str, object]) -> str: thread_ids = sorted( normalize_pst_chat_thread_id(value).lower() for value in cluster["thread_ids"] if normalize_pst_chat_thread_id(value) ) if thread_ids: return f"thread:{thread_ids[0]}" source_folder_paths = sorted( str(value).lower() for value in cluster["source_folder_paths"] if normalize_whitespace(str(value)) ) if source_folder_paths: return f"folder:{source_folder_paths[0]}" source_item_ids = sorted( normalize_whitespace(str(document.get("source_item_id") or "")).lower() for document in list(cluster["documents"]) if normalize_whitespace(str(document.get("source_item_id") or "")) ) if source_item_ids: return f"item:{source_item_ids[0]}" return f"doc:{int(cluster['documents'][0]['id'])}" def pst_chat_conversation_source_locator(cluster: dict[str, object]) -> str: if any(normalize_pst_chat_thread_id(value) for value in cluster["thread_ids"]): return filesystem_dataset_locator() return next( ( str(value) for value in sorted(cluster["source_rel_paths"]) if normalize_whitespace(str(value)) ), None, ) or filesystem_dataset_locator() def derive_pst_chat_conversation_display_name(cluster: dict[str, object]) -> str: generic_folder_names = {"conversation history", "teamsmessagesdata", "teamsmeetings"} participant_summary = render_display_name_title(list(cluster["participant_names"]), max_names=4) if "chat" in cluster["thread_types"] and participant_summary: return participant_summary source_folder_paths = sorted( str(value) for value in cluster["source_folder_paths"] if normalize_whitespace(str(value)) ) for folder_path in source_folder_paths: leaf_name = normalize_whitespace(str(folder_path).split("/")[-1]) if leaf_name and leaf_name.lower() not in generic_folder_names: return leaf_name if participant_summary: return participant_summary for document in list(cluster["documents"]): title = normalize_whitespace(str(document.get("title") or "")) if title: return title for folder_path in source_folder_paths: leaf_name = normalize_whitespace(str(folder_path).split("/")[-1]) if leaf_name: return leaf_name for participants in sorted(str(value) for value in cluster["participants"] if normalize_whitespace(str(value))): if participants: return participants return "Chat conversation" def assign_pst_chat_conversations(connection: sqlite3.Connection) -> dict[str, int]: documents = list_pst_chat_conversation_documents(connection) if not documents: return { "pst_chat_conversations": 0, "pst_chat_documents_reassigned": 0, "pst_chat_child_documents_updated": 0, } clusters: list[dict[str, object]] = [] manual_cluster_ids_by_conversation_id: dict[int, int] = {} auto_cluster_ids_by_scope: dict[tuple[str, str], int] = {} for document in documents: if document["assignment_mode"] != CONVERSATION_ASSIGNMENT_MODE_MANUAL: continue existing_conversation_id = document.get("existing_conversation_id") if existing_conversation_id is None: continue cluster_id = manual_cluster_ids_by_conversation_id.get(int(existing_conversation_id)) if cluster_id is None: cluster_id = len(clusters) manual_cluster_ids_by_conversation_id[int(existing_conversation_id)] = cluster_id clusters.append(create_pst_chat_conversation_cluster(manual_conversation_id=int(existing_conversation_id))) add_document_to_pst_chat_cluster(clusters[cluster_id], document) auto_documents = [document for document in documents if document["assignment_mode"] != CONVERSATION_ASSIGNMENT_MODE_MANUAL] for document in auto_documents: scope_key = pst_chat_cluster_scope_key(document) cluster_id = auto_cluster_ids_by_scope.get(scope_key) if cluster_id is None: cluster_id = len(clusters) auto_cluster_ids_by_scope[scope_key] = cluster_id clusters.append(create_pst_chat_conversation_cluster()) add_document_to_pst_chat_cluster(clusters[cluster_id], document) conversation_id_by_document_id: dict[int, int] = {} unique_conversation_ids: set[int] = set() for cluster in clusters: manual_conversation_id = cluster["manual_conversation_id"] if manual_conversation_id is not None: conversation_id = int(manual_conversation_id) else: conversation_id = upsert_conversation_row( connection, source_kind=PST_SOURCE_KIND, source_locator=pst_chat_conversation_source_locator(cluster), conversation_key=derive_pst_chat_conversation_key(cluster), conversation_type="chat", display_name=derive_pst_chat_conversation_display_name(cluster), ) unique_conversation_ids.add(conversation_id) for document in list(cluster["documents"]): conversation_id_by_document_id[int(document["id"])] = conversation_id documents_reassigned = 0 now = utc_now() for document in auto_documents: target_conversation_id = conversation_id_by_document_id.get(int(document["id"])) if target_conversation_id is None: continue if ( document.get("existing_conversation_id") == target_conversation_id and document.get("assignment_mode") == CONVERSATION_ASSIGNMENT_MODE_AUTO ): continue connection.execute( """ UPDATE documents SET conversation_id = ?, conversation_assignment_mode = ?, updated_at = ? WHERE id = ? AND COALESCE(conversation_assignment_mode, ?) != ? """, ( target_conversation_id, CONVERSATION_ASSIGNMENT_MODE_AUTO, now, int(document["id"]), CONVERSATION_ASSIGNMENT_MODE_AUTO, CONVERSATION_ASSIGNMENT_MODE_MANUAL, ), ) documents_reassigned += 1 child_documents_updated = sync_child_document_conversations( connection, parent_content_types={"Chat"}, parent_source_kinds={PST_SOURCE_KIND}, ) return { "pst_chat_conversations": len(unique_conversation_ids), "pst_chat_documents_reassigned": documents_reassigned, "pst_chat_child_documents_updated": child_documents_updated, } def clear_unsupported_conversation_assignments(connection: sqlite3.Connection) -> int: cursor = connection.execute( """ UPDATE documents SET conversation_id = NULL, conversation_assignment_mode = ?, updated_at = ? WHERE parent_document_id IS NULL AND conversation_id IS NOT NULL AND lifecycle_status NOT IN ('missing', 'deleted') AND LOWER(TRIM(COALESCE(content_type, ''))) NOT IN ('email', 'chat') """, (CONVERSATION_ASSIGNMENT_MODE_AUTO, utc_now()), ) return int(cursor.rowcount or 0) def assign_supported_conversations(connection: sqlite3.Connection) -> dict[str, int]: clear_unsupported_conversation_assignments(connection) email_assignment = assign_email_conversations(connection) pst_chat_assignment = assign_pst_chat_conversations(connection) return { **email_assignment, **pst_chat_assignment, } def list_active_conversation_ids(connection: sqlite3.Connection) -> list[int]: rows = connection.execute( """ SELECT DISTINCT conversation_id FROM documents WHERE conversation_id IS NOT NULL AND lifecycle_status NOT IN ('missing', 'deleted') AND COALESCE(child_document_kind, '') != ? ORDER BY conversation_id ASC """, (CHILD_DOCUMENT_KIND_ATTACHMENT,), ).fetchall() return [int(row["conversation_id"]) for row in rows if row["conversation_id"] is not None] def conversation_preview_date_hint(value: object) -> str | None: candidate = Path(str(value or "")).stem if not candidate: return None try: return f"{date.fromisoformat(candidate).isoformat()}T00:00:00Z" except ValueError: return None def conversation_preview_primary_timestamp(document: dict[str, object]) -> str | None: return ( normalize_datetime(document.get("date_created")) or normalize_datetime(document.get("date_modified")) or conversation_preview_date_hint(document.get("source_rel_path")) or conversation_preview_date_hint(document.get("rel_path")) ) def conversation_preview_sort_key(document: dict[str, object]) -> tuple[str, int, str, int]: return ( conversation_preview_primary_timestamp(document) or "9999-12-31T23:59:59Z", 0 if document.get("parent_document_id") is None else 1, normalize_whitespace(str(document.get("control_number") or "")), int(document["id"]), ) def conversation_preview_segment_mode(conversation_type: object, documents: list[dict[str, object]]) -> str: if normalize_whitespace(str(conversation_type or "")).lower() != "email": return "monthly" total_chars = sum(len(str(document.get("text_content") or "")) for document in documents) if total_chars > CONVERSATION_PREVIEW_MAX_CHARS: return "yearly" return "single" def conversation_preview_segment_key(document: dict[str, object], *, segment_mode: str) -> str: if segment_mode == "single": return "all" timestamp = conversation_preview_primary_timestamp(document) if not timestamp: return "undated" if segment_mode == "yearly": return timestamp[:4] return timestamp[:7] def conversation_preview_segment_label(segment_key: str, *, segment_mode: str) -> str: if segment_key == "all": return "All messages" if segment_key == "undated": return "Undated" if segment_mode == "yearly": return segment_key try: return date.fromisoformat(f"{segment_key}-01").strftime("%B %Y") except ValueError: return segment_key def conversation_documents_are_chat(documents: list[dict[str, object]]) -> bool: return bool(documents) and all( normalize_whitespace(str(document.get("content_type") or "")).lower() == "chat" for document in documents ) def conversation_preview_writes_aggregate_artifacts( conversation_type: object, documents: list[dict[str, object]], *, segment_mode: str, ) -> bool: _ = segment_mode if ( normalize_whitespace(str(conversation_type or "")).lower() == "chat" or conversation_documents_are_chat(documents) ): return bool(documents) return len(documents) > 1 def conversation_preview_document_heading(document: dict[str, object]) -> str: for value in ( document.get("title"), document.get("subject"), document.get("file_name"), document.get("control_number"), ): normalized = normalize_whitespace(str(value or "")) if normalized: return normalized return f"Document {int(document['id'])}" def conversation_preview_document_kind(document: dict[str, object]) -> str: child_kind = normalize_whitespace(str(document.get("child_document_kind") or "")).lower() content_type = normalize_whitespace(str(document.get("content_type") or "Document")) if child_kind == CHILD_DOCUMENT_KIND_REPLY_THREAD: return "Reply thread" if content_type == "Email": return "Email message" if content_type == "Chat": return "Conversation document" return content_type or "Document" def conversation_preview_participants(documents: list[dict[str, object]]) -> str | None: participants: list[str] = [] seen: set[str] = set() for document in documents: append_unique_participants( participants, seen, [ normalize_whitespace(str(document.get("participants") or "")) or None, normalize_whitespace(str(document.get("author") or "")) or None, normalize_whitespace(str(document.get("recipients") or "")) or None, ], ) return ", ".join(participants) or None def conversation_preview_bounds(documents: list[dict[str, object]]) -> tuple[str | None, str | None]: timestamps = sorted( timestamp for timestamp in ( conversation_preview_primary_timestamp(document) for document in documents ) if timestamp ) if not timestamps: return None, None return timestamps[0], timestamps[-1] EMAIL_PREVIEW_BODY_SOURCE_PATTERN = re.compile( r"]*data-retriever-email-body-source\b[^>]*>(.*?)", flags=re.IGNORECASE | re.DOTALL, ) EMAIL_PREVIEW_QUOTED_DETAILS_PATTERN = re.compile( r'(?is)]*class\s*=\s*(["\'])[^"\']*\bgmail-message-quoted\b[^"\']*\1[^>]*>.*?' ) EMAIL_QUOTED_REPLY_SEPARATOR_PATTERN = re.compile( r"^(?:On .+ wrote:|Begin forwarded message:|-{2,}\s*Original Message\s*-{2,}|-{2,}\s*Forwarded message\s*-{2,})$", flags=re.IGNORECASE, ) EMAIL_QUOTED_REPLY_HEADER_PATTERN = re.compile( r"^(?:From|Sent|To|Cc|Bcc|Subject|Date):\s+", flags=re.IGNORECASE, ) EMAIL_HTML_QUOTED_REPLY_START_PATTERNS = ( re.compile( r'(?is)<(?:div|section|table|blockquote)\b[^>]*class\s*=\s*(["\'])[^"\']*\bgmail_(?:quote|attr)\b[^"\']*\1' ), re.compile(r'(?is)]*type\s*=\s*(["\'])cite\1[^>]*>'), re.compile( r"(?is)<(?:div|p|span|font|td)\b[^>]*>\s*" r"(?:On .+ wrote:|Begin forwarded message:|-{2,}\s*(?:Original|Forwarded) Message\s*-{2,})" ), re.compile(r"(?is) None: super().__init__(convert_charrefs=False) self.require_selected_card = require_selected_card self.card_depth: int | None = None self.body_depth: int | None = None self.parts: list[str] = [] self.extracted_html: str | None = None @staticmethod def classes(attrs: list[tuple[str, str | None]]) -> set[str]: class_value = next( ( str(value or "") for key, value in attrs if normalize_whitespace(str(key or "")).lower() == "class" ), "", ) return {token for token in normalize_whitespace(class_value).split(" ") if token} def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: if self.extracted_html is not None: return if self.body_depth is not None: self.parts.append(self.get_starttag_text()) self.body_depth += 1 return classes = self.classes(attrs) if self.card_depth is None: if tag == "article" and "gmail-message-card" in classes and ( not self.require_selected_card or "gmail-message-card--selected" in classes ): self.card_depth = 1 return self.card_depth += 1 if tag == "div" and "gmail-message-body" in classes: self.body_depth = 0 def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: if self.body_depth is not None: self.parts.append(self.get_starttag_text()) def handle_endtag(self, tag: str) -> None: if self.extracted_html is None and self.body_depth is not None: if self.body_depth == 0: extracted = "".join(self.parts).strip() self.extracted_html = extracted or None self.parts.clear() self.body_depth = None else: self.parts.append(f"") self.body_depth -= 1 if self.card_depth is not None: self.card_depth -= 1 if self.card_depth <= 0: self.card_depth = None def handle_data(self, data: str) -> None: if self.body_depth is not None: self.parts.append(data) def handle_entityref(self, name: str) -> None: if self.body_depth is not None: self.parts.append(f"&{name};") def handle_charref(self, name: str) -> None: if self.body_depth is not None: self.parts.append(f"&#{name};") def handle_comment(self, data: str) -> None: if self.body_depth is not None: self.parts.append(f"") def handle_decl(self, decl: str) -> None: if self.body_depth is not None: self.parts.append(f"") def handle_pi(self, data: str) -> None: if self.body_depth is not None: self.parts.append(f"") def extract_visible_email_preview_body_html(preview_html: str) -> str | None: for require_selected_card in (True, False): parser = EmailPreviewBodyHTMLExtractor(require_selected_card=require_selected_card) parser.feed(preview_html) parser.close() if parser.extracted_html: return parser.extracted_html return None def build_email_preview_head_html() -> str: return ( "" ) def email_preview_document_id(document: dict[str, object]) -> int | None: try: raw_value = document.get("id") if raw_value is None: return None value = int(raw_value) return value if value > 0 else None except (TypeError, ValueError): return None def normalize_email_preview_address(value: object) -> tuple[str | None, str | None]: normalized = normalize_whitespace(str(value or "")) if not normalized: return None, None addresses = getaddresses([normalized]) if addresses: name, address = addresses[0] normalized_name = normalize_whitespace(name) normalized_address = normalize_whitespace(address) if normalized_name or normalized_address: return normalized_name or normalized_address or None, normalized_address or None return normalized, None def format_email_preview_person(value: object) -> str | None: name, address = normalize_email_preview_address(value) if name and address and name != address: return f"{name} <{address}>" return name or address def summarize_email_preview_recipients(value: object) -> str | None: normalized = normalize_whitespace(str(value or "")) if not normalized: return None addresses = getaddresses([normalized]) formatted_addresses = [ formatted for formatted in (format_email_preview_person(f"{name} <{address}>") if address else normalize_whitespace(name) for name, address in addresses) if formatted ] if formatted_addresses: return f"to {', '.join(formatted_addresses)}" return f"to {normalized}" def split_email_preview_text_content(text_content: str) -> tuple[str, str | None]: normalized_text = text_content.replace("\r\n", "\n").replace("\r", "\n").strip("\n") if not normalized_text: return "", None lines = normalized_text.split("\n") has_visible_content = False for index, raw_line in enumerate(lines): stripped = raw_line.strip() if not stripped: continue if has_visible_content: if raw_line.lstrip().startswith(">") or EMAIL_QUOTED_REPLY_SEPARATOR_PATTERN.match(stripped): visible_text = "\n".join(lines[:index]).strip() quoted_text = "\n".join(lines[index:]).strip() if visible_text and quoted_text: return visible_text, quoted_text if index > 0 and not lines[index - 1].strip() and EMAIL_QUOTED_REPLY_HEADER_PATTERN.match(stripped): visible_text = "\n".join(lines[:index]).strip() quoted_text = "\n".join(lines[index:]).strip() if visible_text and quoted_text: return visible_text, quoted_text has_visible_content = True return normalized_text, None def render_email_preview_text_html(text: str) -> str: normalized_text = text.replace("\r\n", "\n").replace("\r", "\n").strip("\n") if not normalized_text: return "" return "
".join(html.escape(line) for line in normalized_text.split("\n")) def split_email_preview_visible_text_fragments(text: str, *, max_fragments: int) -> list[str]: normalized_text = text.replace("\r\n", "\n").replace("\r", "\n").strip() if not normalized_text: return [] if max_fragments <= 1: return [normalized_text] fragments = [ fragment.strip() for fragment in re.split(r"\n\s*\n", normalized_text) if fragment.strip() ] if not fragments: return [normalized_text] if len(fragments) <= max_fragments: return fragments return [*fragments[: max_fragments - 1], "\n\n".join(fragments[max_fragments - 1 :])] class EmailPreviewVisibleTextNodeCounter(HTMLParser): def __init__(self) -> None: super().__init__(convert_charrefs=True) self.visible_text_node_count = 0 self.literal_tag_stack: list[str] = [] def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: del attrs if tag.lower() in {"script", "style"}: self.literal_tag_stack.append(tag.lower()) def handle_endtag(self, tag: str) -> None: if self.literal_tag_stack and self.literal_tag_stack[-1] == tag.lower(): self.literal_tag_stack.pop() def handle_data(self, data: str) -> None: if self.literal_tag_stack: return if normalize_whitespace(data): self.visible_text_node_count += 1 def count_email_preview_visible_text_nodes(html_content: str | None) -> int: normalized_html = str(html_content or "").strip() if not normalized_html: return 0 parser = EmailPreviewVisibleTextNodeCounter() parser.feed(normalized_html) parser.close() return parser.visible_text_node_count class EmailPreviewActiveTextHTMLRewriter(HTMLParser): def __init__( self, *, fragments: list[str], visible_text_node_count: int, ) -> None: super().__init__(convert_charrefs=True) self.fragments = list(fragments) self.visible_text_node_count = visible_text_node_count self.processed_visible_text_nodes = 0 self.literal_tag_stack: list[str] = [] self.parts: list[str] = [] def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: del attrs self.parts.append(self.get_starttag_text()) if tag.lower() in {"script", "style"}: self.literal_tag_stack.append(tag.lower()) def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: del tag, attrs self.parts.append(self.get_starttag_text()) def handle_endtag(self, tag: str) -> None: self.parts.append(f"") if self.literal_tag_stack and self.literal_tag_stack[-1] == tag.lower(): self.literal_tag_stack.pop() def handle_data(self, data: str) -> None: if self.literal_tag_stack: self.parts.append(data) return if not normalize_whitespace(data): self.parts.append(data) return self.processed_visible_text_nodes += 1 if not self.fragments: return fragment = self.fragments.pop(0) if self.processed_visible_text_nodes == self.visible_text_node_count and self.fragments: fragment = "\n\n".join([fragment, *self.fragments]) self.fragments.clear() self.parts.append(render_email_preview_text_html(fragment)) def handle_comment(self, data: str) -> None: self.parts.append(f"") def handle_decl(self, decl: str) -> None: self.parts.append(f"") def handle_pi(self, data: str) -> None: self.parts.append(f"") def rewrite_email_body_html_with_active_text( html_content: str | None, active_text: str, *, strip_quoted_history: bool = False, ) -> str | None: normalized_html = str(html_content or "").strip() if not normalized_html: return None source_html = ( strip_email_reply_history_html(normalized_html) if strip_quoted_history else normalized_html ) if not source_html: return None visible_text_node_count = count_email_preview_visible_text_nodes(source_html) if visible_text_node_count <= 0: return None fragments = split_email_preview_visible_text_fragments( active_text, max_fragments=visible_text_node_count, ) if not fragments: return None parser = EmailPreviewActiveTextHTMLRewriter( fragments=fragments, visible_text_node_count=visible_text_node_count, ) parser.feed(source_html) parser.close() rewritten_html = "".join(parser.parts).strip() return rewritten_html or None def email_html_has_visible_content(html_content: str | None) -> bool: return bool(normalize_whitespace(strip_html_tags(str(html_content or "")))) def strip_email_reply_history_html(html_content: str | None) -> str | None: normalized_html = str(html_content or "").strip() if not normalized_html: return None split_index = min( ( match.start() for pattern in EMAIL_HTML_QUOTED_REPLY_START_PATTERNS if (match := pattern.search(normalized_html)) is not None ), default=None, ) if split_index is None: return normalized_html candidate = re.sub(r"(?is)(?:|\s| )+$", "", normalized_html[:split_index]).strip() return candidate if email_html_has_visible_content(candidate) else None HTML_PREVIEW_RELATIVE_ASSET_ATTR_PATTERN = re.compile( r"(?is)(\s(?:src|poster)\s*=\s*)([\"'])(.*?)\2" ) ABSOLUTE_OR_SPECIAL_URL_PATTERN = re.compile(r"^[a-zA-Z][a-zA-Z0-9+.-]*:") def preview_asset_url_is_relative(value: str) -> bool: normalized = normalize_whitespace(str(value or "")) if not normalized: return False return not ( normalized.startswith(("#", "/", "//")) or ABSOLUTE_OR_SPECIAL_URL_PATTERN.match(normalized) is not None ) def preview_relative_asset_rel_path( value: str, *, source_preview_rel_path: object, ) -> str | None: normalized_value = str(value or "") if not preview_asset_url_is_relative(normalized_value): return None source_rel_path = normalize_internal_rel_path(Path(str(source_preview_rel_path or ""))) if not source_rel_path: return None source_dir = posixpath.dirname(source_rel_path) split_at = len(normalized_value) for separator in ("?", "#"): index = normalized_value.find(separator) if index >= 0: split_at = min(split_at, index) url_path = normalized_value[:split_at] if not url_path: return None asset_rel_path = posixpath.normpath( posixpath.join(source_dir, urllib_request.url2pathname(url_path)) ) if asset_rel_path == ".." or asset_rel_path.startswith("../"): return None return asset_rel_path def preview_relative_asset_rel_paths_for_html_body( html_body: str, *, source_preview_rel_path: object, ) -> list[str]: if not html_body or not source_preview_rel_path: return [] rel_paths: list[str] = [] seen_paths: set[str] = set() for match in HTML_PREVIEW_RELATIVE_ASSET_ATTR_PATTERN.finditer(html_body): asset_rel_path = preview_relative_asset_rel_path( match.group(3), source_preview_rel_path=source_preview_rel_path, ) if asset_rel_path is None or asset_rel_path in seen_paths: continue seen_paths.add(asset_rel_path) rel_paths.append(asset_rel_path) return rel_paths def rebase_preview_relative_asset_url( value: str, *, source_preview_rel_path: object, target_preview_rel_path: object, ) -> str: normalized_value = str(value or "") if not preview_asset_url_is_relative(normalized_value): return normalized_value target_rel_path = normalize_internal_rel_path(Path(str(target_preview_rel_path or ""))) asset_rel_path = preview_relative_asset_rel_path( normalized_value, source_preview_rel_path=source_preview_rel_path, ) if asset_rel_path is None or not target_rel_path: return normalized_value target_dir = posixpath.dirname(target_rel_path) split_at = len(normalized_value) for separator in ("?", "#"): index = normalized_value.find(separator) if index >= 0: split_at = min(split_at, index) suffix = normalized_value[split_at:] rebased_path = posixpath.relpath(asset_rel_path, start=target_dir or ".") return urllib_request.pathname2url(rebased_path) + suffix def rebase_preview_relative_asset_paths( html_body: str, *, source_preview_rel_path: object, target_preview_rel_path: object, ) -> str: if not html_body or not source_preview_rel_path or not target_preview_rel_path: return html_body def replace_attr(match: re.Match[str]) -> str: rebased_url = rebase_preview_relative_asset_url( match.group(3), source_preview_rel_path=source_preview_rel_path, target_preview_rel_path=target_preview_rel_path, ) return f"{match.group(1)}{match.group(2)}{html.escape(rebased_url, quote=True)}{match.group(2)}" return HTML_PREVIEW_RELATIVE_ASSET_ATTR_PATTERN.sub(replace_attr, html_body) def preview_body_html_for_output( document: dict[str, object], body_html: str | None, *, target_preview_rel_path: object = None, ) -> str | None: if not isinstance(body_html, str) or not body_html.strip(): return None if target_preview_rel_path is None: return body_html return rebase_preview_relative_asset_paths( body_html, source_preview_rel_path=document.get("standalone_preview_rel_path"), target_preview_rel_path=target_preview_rel_path, ) def email_preview_body_html_has_rendered_wrappers(html_body: str) -> bool: return any( token in html_body for token in ( 'class="gmail-message-rendered-html"', 'class="gmail-message-plain"', 'class="gmail-message-quoted"', ) ) def build_email_message_body_content_html( document: dict[str, object], *, body_html: str | None = None, target_preview_rel_path: object = None, strip_quoted_history: bool = False, ) -> str: text_content, calendar_invites = extract_calendar_invites_from_text_content( str(document.get("text_content") or "") ) visible_text, quoted_text = split_email_preview_text_content(text_content) calendar_invites_html = render_html_preview_calendar_invite_cards(calendar_invites) preferred_body_html = body_html if preferred_body_html is None: stored_body_html = document.get("standalone_preview_body_html") preferred_body_html = ( str(stored_body_html) if isinstance(stored_body_html, str) and stored_body_html.strip() else None ) preferred_body_html = preview_body_html_for_output( document, preferred_body_html, target_preview_rel_path=target_preview_rel_path, ) if ( preferred_body_html and text_content and document_active_text_should_render_above_preserved_email_body(document) ): normalized_html = HTML_PREVIEW_CALENDAR_INVITES_PATTERN.sub("", preferred_body_html).strip() normalized_html = EMAIL_PREVIEW_QUOTED_DETAILS_PATTERN.sub("", normalized_html).strip() if strip_quoted_history: stripped_html = strip_email_reply_history_html(normalized_html) if stripped_html is not None: normalized_html = stripped_html body_parts: list[str] = [] if calendar_invites_html: body_parts.append(calendar_invites_html) body_parts.append( '
' f"{render_email_preview_text_html(text_content)}" "
" ) if normalized_html: if email_preview_body_html_has_rendered_wrappers(normalized_html): body_parts.append(normalized_html) else: body_parts.append(f'
{normalized_html}
') return "".join(body_parts) if preferred_body_html and document_active_text_differs_from_source(document): normalized_html = HTML_PREVIEW_CALENDAR_INVITES_PATTERN.sub("", preferred_body_html).strip() normalized_html = EMAIL_PREVIEW_QUOTED_DETAILS_PATTERN.sub("", normalized_html).strip() visible_text_node_count = count_email_preview_visible_text_nodes( strip_email_reply_history_html(normalized_html) if strip_quoted_history else normalized_html ) rewritten_html = rewrite_email_body_html_with_active_text( normalized_html, visible_text, strip_quoted_history=strip_quoted_history, ) body_parts: list[str] = [] if calendar_invites_html: body_parts.append(calendar_invites_html) if rewritten_html: if email_preview_body_html_has_rendered_wrappers(rewritten_html): body_parts.append(rewritten_html) else: body_parts.append(f'
{rewritten_html}
') elif visible_text: body_parts.append( f'
{render_email_preview_text_html(visible_text)}
' ) if visible_text_node_count <= 0 and normalized_html: if email_preview_body_html_has_rendered_wrappers(normalized_html): body_parts.append(normalized_html) else: body_parts.append(f'
{normalized_html}
') elif visible_text_node_count <= 0 and normalized_html: if email_preview_body_html_has_rendered_wrappers(normalized_html): body_parts.append(normalized_html) else: body_parts.append(f'
{normalized_html}
') elif not body_parts: body_parts.append( '
No extracted text available.
' ) if quoted_text and not strip_quoted_history: body_parts.append( "
" "Quoted text" f"
{html.escape(quoted_text)}
" "
" ) return "".join(body_parts) if preferred_body_html: normalized_html = preferred_body_html.strip() if strip_quoted_history: stripped_html = strip_email_reply_history_html(normalized_html) if stripped_html is not None: normalized_html = stripped_html elif quoted_text and visible_text: normalized_html = "" normalized_html = HTML_PREVIEW_CALENDAR_INVITES_PATTERN.sub("", normalized_html).strip() if normalized_html and email_preview_body_html_has_rendered_wrappers(normalized_html): return f"{calendar_invites_html}{normalized_html}" if calendar_invites_html else normalized_html normalized_html = re.sub(r"(?is)]*>\s*", "", normalized_html).strip() if normalized_html: rendered_html = f'
{normalized_html}
' return f"{calendar_invites_html}{rendered_html}" if calendar_invites_html else rendered_html body_parts: list[str] = [] if calendar_invites_html: body_parts.append(calendar_invites_html) plain_text = visible_text or ("No extracted text available." if not body_parts else "") if plain_text: body_parts.append( f'
{render_email_preview_text_html(plain_text)}
' ) elif not body_parts: body_parts.append( '
No extracted text available.
' ) if quoted_text and not strip_quoted_history: body_parts.append( "
" "Quoted text" f"
{html.escape(quoted_text)}
" "
" ) return "".join(body_parts) def build_email_message_card_html( document: dict[str, object], *, body_html: str | None = None, selected: bool = False, attachment_links: list[dict[str, str]] | None = None, target_preview_rel_path: object = None, strip_quoted_history: bool = False, ) -> str: author_name, author_email = normalize_email_preview_address(document.get("author")) author_label = author_name or author_email or "Unknown sender" author_email_html = ( f'<{html.escape(author_email)}>' if author_email and author_email != author_label else "" ) recipients_line = summarize_email_preview_recipients(document.get("recipients")) timestamp_label = format_chat_preview_timestamp(conversation_preview_primary_timestamp(document)) or "" avatar_background, avatar_foreground = chat_avatar_colors(author_label) avatar_html = build_chat_avatar_svg( chat_avatar_initials(author_label), avatar_background, avatar_foreground, author_label, ) message_body_html = build_email_message_body_content_html( document, body_html=body_html, target_preview_rel_path=target_preview_rel_path, strip_quoted_history=strip_quoted_history, ) attachment_links_html = render_html_preview_attachment_links(attachment_links or []) card_classes = "gmail-message-card gmail-message-card--selected" if selected else "gmail-message-card" anchor = conversation_preview_anchor(document_id) if (document_id := email_preview_document_id(document)) is not None else None anchor_attr = f' id="{html.escape(anchor)}"' if anchor else "" recipient_line_html = ( f'
{html.escape(recipients_line)}
' if recipients_line else "" ) time_html = f'
{html.escape(timestamp_label)}
' if timestamp_label else "" return ( f'
' f"{avatar_html}" '
' '
' '
' '
' f'{html.escape(author_label)}' f"{author_email_html}" "
" f"{recipient_line_html}" "
" f"{time_html}" "
" f'
{message_body_html}
' f"{attachment_links_html}" "
" "
" ) def build_email_thread_summary_html( documents: list[dict[str, object]], *, summary_documents: list[dict[str, object]] | None = None, position_index: int | None = None, segment_label: str | None = None, segment_count: int | None = None, ) -> str: summary_scope = summary_documents if summary_documents is not None else documents if not summary_scope: return "" total_messages = len(summary_scope) if total_messages <= 1 and not (segment_label and (segment_count or 0) > 1): return "" started_at, last_message_at = conversation_preview_bounds(summary_scope) participants = conversation_preview_participants(summary_scope) pills: list[tuple[str, bool]] = [] if segment_label and (segment_count or 0) > 1: pills.append((segment_label, False)) if total_messages > 1: pills.append((f"{total_messages} messages", False)) started_label = format_chat_preview_timestamp(started_at) or started_at or "" if started_label: pills.append((f"Created {started_label}", False)) last_message_label = format_chat_preview_timestamp(last_message_at) or last_message_at or "" if last_message_label and last_message_label != started_label: pills.append((f"Last modified {last_message_label}", False)) if participants: pills.append((f"Participants {participants}", False)) if position_index is not None and total_messages > 1: pills.append((f"Viewing message {position_index} of {total_messages}", True)) if not pills: return "" pills_html = "".join( f'{html.escape(label)}' for label, is_active in pills if label ) return f'
{pills_html}
' if pills_html else "" def build_email_thread_title_html( thread_title: str, *, thread_link_href: str | None = None, thread_position_label: str | None = None, ) -> str: if not thread_link_href and not thread_position_label: return html.escape(thread_title) title_html = ( f'{html.escape(thread_title)}' if thread_link_href else f'{html.escape(thread_title)}' ) if thread_position_label: title_html += f'({html.escape(thread_position_label)})' return title_html def build_email_thread_preview_html( *, thread_title: str, documents: list[dict[str, object]], summary_documents: list[dict[str, object]] | None = None, page_title: str, selected_document_id: int | None = None, position_index: int | None = None, segment_label: str | None = None, segment_count: int | None = None, attachment_links_by_document_id: dict[int, list[dict[str, str]]] | None = None, body_source_document: dict[str, object] | None = None, body_source_html: str | None = None, thread_link_href: str | None = None, thread_position_label: str | None = None, target_preview_rel_path: object = None, strip_quoted_history: bool = False, newest_first: bool = False, ) -> str: selected_card_script = "" if selected_document_id is not None: selected_card_script = ( "" ) summary_html = build_email_thread_summary_html( documents, summary_documents=summary_documents, position_index=position_index, segment_label=segment_label, segment_count=segment_count, ) title_html = build_email_thread_title_html( thread_title, thread_link_href=thread_link_href, thread_position_label=thread_position_label, ) body_source_document_id = ( email_preview_document_id(body_source_document) if body_source_document is not None else None ) render_documents = list(reversed(documents)) if newest_first else list(documents) message_cards = [] for document in render_documents: document_id = email_preview_document_id(document) message_cards.append( build_email_message_card_html( document, body_html=( body_source_html if body_source_document_id is not None and document_id == body_source_document_id else None ), selected=(selected_document_id is not None and document_id == selected_document_id), attachment_links=( (attachment_links_by_document_id or {}).get(document_id, []) if document_id is not None else None ), target_preview_rel_path=target_preview_rel_path, strip_quoted_history=strip_quoted_history, ) ) header_kicker = ( f'

{html.escape(segment_label)}

' if segment_label and (segment_count or 0) > 1 else "" ) return ( "" "" '' f"{html.escape(page_title)}" f"{build_email_preview_head_html()}" "" '
' '
' f"{header_kicker}" f'

{title_html}

' f"{summary_html}" "
" f'
{"".join(message_cards)}
' "
" f"{selected_card_script}" "" ) def build_email_message_preview_html( document: dict[str, object], *, body_html: str | None, conversation_row: sqlite3.Row | None = None, conversation_documents: list[dict[str, object]] | None = None, position_index: int | None = None, thread_link_href: str | None = None, attachment_links: list[dict[str, str]] | None = None, target_preview_rel_path: object = None, ) -> str: document_title = conversation_preview_document_heading(document) or "Retriever Email Preview" document_id = email_preview_document_id(document) attachment_links_by_document_id = ( {document_id: attachment_links} if document_id is not None and attachment_links else None ) if conversation_row is not None and conversation_documents is not None and len(conversation_documents) > 1: email_documents = [ item for item in conversation_documents if normalize_whitespace(str(item.get("content_type") or "")).lower() == "email" ] selected_position = position_index if selected_position is None and document_id is not None: for index, candidate in enumerate(email_documents, start=1): if email_preview_document_id(candidate) == document_id: selected_position = index break if selected_position is None: selected_rel_path = normalize_whitespace(str(document.get("rel_path") or "")) for index, candidate in enumerate(email_documents, start=1): if normalize_whitespace(str(candidate.get("rel_path") or "")) == selected_rel_path: selected_position = index break timeline_documents = ( email_documents[:selected_position] if selected_position is not None and selected_position > 0 else [document] ) thread_title = ( normalize_whitespace(str(conversation_row["display_name"] or "")) or normalize_generated_document_title(document.get("subject") or document.get("title")) or document_title ) thread_message_count = len(email_documents) thread_position_label = ( f"{selected_position}/{thread_message_count} in thread" if selected_position is not None and thread_message_count > 1 else None ) return build_email_thread_preview_html( thread_title=thread_title, documents=timeline_documents, summary_documents=email_documents, page_title=( f"{thread_title} ({thread_position_label})" if thread_position_label else thread_title ), selected_document_id=document_id, position_index=selected_position, body_source_document=document, body_source_html=body_html, thread_link_href=thread_link_href, thread_position_label=thread_position_label, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=target_preview_rel_path, strip_quoted_history=True, newest_first=True, ) thread_title = normalize_generated_document_title(document.get("subject") or document.get("title")) or document_title return build_email_thread_preview_html( thread_title=thread_title, documents=[document], page_title=document_title, selected_document_id=document_id, body_source_document=document, body_source_html=body_html, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=target_preview_rel_path, ) def default_email_message_preview_rel_path(document: dict[str, object]) -> str | None: rel_path = normalize_whitespace(str(document.get("rel_path") or "")) if not rel_path: return None preview_base = preview_base_path_for_rel_path(rel_path) source_kind = normalize_whitespace(str(document.get("source_kind") or "")).lower() source_item_id = normalize_whitespace(str(document.get("source_item_id") or "")) if source_kind in {PST_SOURCE_KIND, MBOX_SOURCE_KIND} and source_item_id: preview_file_name = container_preview_file_name(source_item_id) else: preview_file_name = f"{Path(rel_path).name}.html" return (preview_base / preview_file_name).as_posix() def rewrite_preserved_email_message_preview( paths: dict[str, Path], *, document: dict[str, object], preview_rows: list[dict[str, object]], conversation_row: sqlite3.Row | None = None, conversation_documents: list[dict[str, object]] | None = None, position_index: int | None = None, thread_rel_path: str | None = None, attachment_links: list[dict[str, str]] | None = None, ) -> None: preferred_rows = sorted( preview_rows, key=lambda row: ( 0 if normalize_whitespace(str(row.get("label") or "")).lower() == "message" else 1, int(row.get("ordinal", 0)), ), ) target_row = next( ( row for row in preferred_rows if normalize_whitespace(str(row.get("preview_type") or "")).lower() == "html" ), None, ) if target_row is None: synthesized_rel_path = default_email_message_preview_rel_path(document) if synthesized_rel_path is None: return target_row = { "rel_preview_path": synthesized_rel_path, "preview_type": "html", "target_fragment": None, "label": "message", "ordinal": 0, } preview_rows.insert(0, target_row) else: preview_rows[:] = [target_row, *[row for row in preview_rows if row is not target_row]] if not normalize_whitespace(str(target_row.get("label") or "")): target_row["label"] = "message" preview_path = paths["state_dir"] / str(target_row["rel_preview_path"]) preview_path.parent.mkdir(parents=True, exist_ok=True) thread_link_href = None if thread_rel_path: thread_path = paths["state_dir"] / thread_rel_path if thread_path.exists(): thread_link_href = relative_preview_href(thread_path, preview_path) body_html = document.get("standalone_preview_body_html") target_preview_rel_path = str(target_row["rel_preview_path"]) preview_path.write_text( build_email_message_preview_html( document, body_html=(str(body_html) if isinstance(body_html, str) and body_html.strip() else None), conversation_row=conversation_row, conversation_documents=conversation_documents, position_index=position_index, thread_link_href=thread_link_href, attachment_links=attachment_links, target_preview_rel_path=target_preview_rel_path, ), encoding="utf-8", ) def extract_standalone_preview_body_html(preview_html: str) -> str | None: cleaned_preview_html = HTML_PREVIEW_ATTACHMENT_LINKS_PATTERN.sub("", preview_html or "") cleaned_preview_html = HTML_PREVIEW_CALENDAR_INVITES_PATTERN.sub("", cleaned_preview_html) visible_body_html = extract_visible_email_preview_body_html(cleaned_preview_html) if visible_body_html is not None: return visible_body_html source_match = EMAIL_PREVIEW_BODY_SOURCE_PATTERN.search(cleaned_preview_html) if source_match is not None: normalized_source = source_match.group(1).strip() return normalized_source or None body_match = re.search( r"]*>(.*)\s*\s*$", cleaned_preview_html, flags=re.IGNORECASE | re.DOTALL, ) if body_match is None: return None body_html = re.sub( r"^\s*]*>.*?\s*", "", body_match.group(1), count=1, flags=re.IGNORECASE | re.DOTALL, ) body_html = re.sub( r"^\s*]*>.*?\s*\s*", "", body_html, count=1, flags=re.IGNORECASE | re.DOTALL, ) normalized = body_html.strip() return normalized or None def load_preserved_preview_rows_by_document_id( connection: sqlite3.Connection, paths: dict[str, Path], document_ids: list[int], ) -> dict[int, list[dict[str, object]]]: if not document_ids: return {} rows = connection.execute( f""" SELECT document_id, rel_preview_path, preview_type, target_fragment, label, ordinal FROM document_previews WHERE document_id IN ({", ".join("?" for _ in document_ids)}) ORDER BY document_id ASC, ordinal ASC, id ASC """, document_ids, ).fetchall() preview_rows_by_document_id: dict[int, list[dict[str, object]]] = defaultdict(list) for row in rows: rel_preview_path = normalize_whitespace(str(row["rel_preview_path"] or "")) if not rel_preview_path or is_conversation_preview_rel_path(rel_preview_path): continue if not (paths["state_dir"] / rel_preview_path).exists(): continue preview_rows_by_document_id[int(row["document_id"])].append( { "rel_preview_path": rel_preview_path, "preview_type": str(row["preview_type"]), "target_fragment": row["target_fragment"], "label": (str(row["label"]) if row["label"] is not None else None), "ordinal": int(row["ordinal"]), } ) return dict(preview_rows_by_document_id) def load_document_preview_body_payload( paths: dict[str, Path], preview_rows: list[dict[str, object]], ) -> dict[str, str] | None: preferred_rows = sorted( preview_rows, key=lambda row: ( 0 if normalize_whitespace(str(row.get("label") or "")).lower() == "message" else 1, int(row.get("ordinal", 0)), ), ) for preview_row in preferred_rows: if normalize_whitespace(str(preview_row.get("preview_type") or "")).lower() != "html": continue preview_path = paths["state_dir"] / str(preview_row["rel_preview_path"]) if not preview_path.exists(): continue body_html = extract_standalone_preview_body_html( preview_path.read_text(encoding="utf-8") ) if body_html: return { "body_html": body_html, "rel_preview_path": str(preview_row["rel_preview_path"]), } return None def load_source_backed_email_preview_body_payload( paths: dict[str, Path], document: dict[str, object], ) -> dict[str, str] | None: rel_path = normalize_whitespace(str(document.get("rel_path") or "")) if not rel_path: return None source_kind = normalize_whitespace(str(document.get("source_kind") or FILESYSTEM_SOURCE_KIND)).lower() if source_kind != FILESYSTEM_SOURCE_KIND: return None source_path = document_absolute_path(paths, rel_path) if not source_path.exists(): return None try: extracted = extract_document(source_path, include_attachments=True) except RetrieverError: return None preview_artifacts = [ artifact for artifact in list(extracted.get("preview_artifacts") or []) if isinstance(artifact, dict) and normalize_whitespace(str(artifact.get("preview_type") or "")).lower() == "html" and normalize_whitespace(str(artifact.get("content") or "")) ] if not preview_artifacts: return None preferred_artifact = sorted( preview_artifacts, key=lambda artifact: ( 0 if normalize_whitespace(str(artifact.get("label") or "")).lower() == "message" else 1, int(artifact.get("ordinal", 0) or 0), ), )[0] body_html = extract_standalone_preview_body_html(str(preferred_artifact["content"])) if not body_html: return None preview_file_name = normalize_whitespace(str(preferred_artifact.get("file_name") or "")) preview_rel_path = ( (preview_base_path_for_rel_path(rel_path) / preview_file_name).as_posix() if preview_file_name else default_email_message_preview_rel_path(document) ) return { "body_html": body_html, "rel_preview_path": preview_rel_path or default_email_message_preview_rel_path(document) or "", } def load_document_preview_body_html( paths: dict[str, Path], preview_rows: list[dict[str, object]], ) -> str | None: payload = load_document_preview_body_payload(paths, preview_rows) return payload["body_html"] if payload else None def rebase_preserved_preview_rows( preview_rows: list[dict[str, object]], *, start_ordinal: int, created_at: str, ) -> list[dict[str, object]]: rebased_rows: list[dict[str, object]] = [] for index, preview_row in enumerate(preview_rows): rebased_rows.append( { "rel_preview_path": str(preview_row["rel_preview_path"]), "preview_type": str(preview_row["preview_type"]), "target_fragment": preview_row.get("target_fragment"), "label": preview_row.get("label"), "ordinal": start_ordinal + index, "created_at": created_at, } ) return rebased_rows def chat_preview_entries_for_document(document: dict[str, object]) -> list[dict[str, object]]: text_content = str(document.get("text_content") or "") parsed_entries = list(iter_chat_transcript_entries(text_content, max_lines=4000)) if parsed_entries: return parsed_entries body = normalize_whitespace(text_content) if not body: return [] speaker = normalize_whitespace(str(document.get("author") or "")) or "Message" timestamp = conversation_preview_primary_timestamp(document) return [ { "speaker": speaker, "body": body, "timestamp": timestamp, "timestamp_label": format_chat_preview_timestamp(timestamp), "avatar_label": chat_avatar_initials(speaker), } ] def chat_preview_headers_for_document(document: dict[str, object]) -> dict[str, str]: date_created = conversation_preview_primary_timestamp(document) date_modified = normalize_datetime(document.get("date_modified")) return { "Author": normalize_whitespace(str(document.get("author") or "")), "Participants": normalize_whitespace(str(document.get("participants") or "")), "Started": format_chat_preview_timestamp(date_created) or date_created or "", "Updated": format_chat_preview_timestamp(date_modified) or date_modified or "", "Title": conversation_preview_document_heading(document), "Messages": str(len(chat_preview_entries_for_document(document)) or 1), } def build_chat_document_preview_html(document: dict[str, object]) -> str: text_content = str(document.get("text_content") or "") return build_chat_preview_html( chat_preview_headers_for_document(document), text_content, document_title=conversation_preview_document_heading(document) or "Chat message", entries=chat_preview_entries_for_document(document), ) def default_chat_document_preview_file_name(document: dict[str, object]) -> str: rel_path = normalize_whitespace(str(document.get("rel_path") or "")) source_kind = normalize_whitespace(str(document.get("source_kind") or "")).lower() source_item_id = normalize_whitespace(str(document.get("source_item_id") or "")) if source_kind in CONTAINER_SOURCE_FILE_TYPES and source_item_id: return container_preview_file_name(source_item_id) if rel_path: return f"{Path(rel_path).name}.html" return "preview.html" def regenerate_chat_preview_for_document( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, ) -> dict[str, object]: documents = load_preview_documents(connection, paths, document_ids=[document_id]) if not documents: return {"status": "skipped", "reason": "unknown_document"} document = dict(documents[0]) if not document_content_type_is_chat(document.get("content_type")): return {"status": "skipped", "reason": "not_chat"} if document.get("conversation_id") is not None: return {"status": "skipped", "reason": "conversation_scope"} rel_path = normalize_whitespace(str(document.get("rel_path") or "")) if not rel_path: return {"status": "skipped", "reason": "missing_rel_path"} preserved_preview_rows = load_preserved_preview_rows_by_document_id( connection, paths, [document_id], ).get(document_id, []) preferred_html_row = next( ( row for row in sorted( preserved_preview_rows, key=lambda item: int(item.get("ordinal", 0)), ) if normalize_whitespace(str(row.get("preview_type") or "")).lower() == "html" ), None, ) preview_file_name = ( Path(str(preferred_html_row["rel_preview_path"])).name if preferred_html_row is not None and preferred_html_row.get("rel_preview_path") else default_chat_document_preview_file_name(document) ) preview_label = ( normalize_whitespace(str(preferred_html_row.get("label") or "")) if preferred_html_row is not None else "" ) preview_artifacts = [ { "file_name": preview_file_name, "preview_type": "html", "label": preview_label or "conversation", "ordinal": 0, "content": build_chat_document_preview_html(document), } ] previous_preview_paths = [ str(preview_row["rel_preview_path"]) for preview_row in connection.execute( "SELECT rel_preview_path FROM document_previews WHERE document_id = ?", (document_id,), ).fetchall() ] preview_rows = write_preview_artifacts(paths, rel_path, preview_artifacts) preserved_non_html_rows = rebase_preserved_preview_rows( [ row for row in preserved_preview_rows if normalize_whitespace(str(row.get("preview_type") or "")).lower() != "html" ], start_ordinal=len(preview_rows), created_at=utc_now(), ) replace_document_preview_rows(connection, document_id, [*preview_rows, *preserved_non_html_rows]) cleanup_unreferenced_preview_files(paths, connection, previous_preview_paths) sync_document_attachment_preview_links(connection, paths, document_id) return {"status": "ok", "preview_rows": len(preview_rows) + len(preserved_non_html_rows)} def regenerate_email_preview_for_document( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, ) -> dict[str, object]: documents = load_preview_documents(connection, paths, document_ids=[document_id]) if not documents: return {"status": "skipped", "reason": "unknown_document"} document = dict(documents[0]) if not document_content_type_is_email(document.get("content_type")): return {"status": "skipped", "reason": "not_email"} if document.get("conversation_id") is not None: return {"status": "skipped", "reason": "conversation_scope"} preserved_preview_rows = load_preserved_preview_rows_by_document_id( connection, paths, [document_id], ).get(document_id, []) previous_preview_paths = [ str(preview_row["rel_preview_path"]) for preview_row in connection.execute( "SELECT rel_preview_path FROM document_previews WHERE document_id = ?", (document_id,), ).fetchall() ] rewrite_preserved_email_message_preview( paths, document=document, preview_rows=preserved_preview_rows, ) preview_rows = rebase_preserved_preview_rows( preserved_preview_rows, start_ordinal=0, created_at=utc_now(), ) replace_document_preview_rows(connection, document_id, preview_rows) cleanup_unreferenced_preview_files(paths, connection, previous_preview_paths) sync_document_attachment_preview_links(connection, paths, document_id) return {"status": "ok", "preview_rows": len(preview_rows)} def build_chat_conversation_preview_html( conversation_row: sqlite3.Row, documents: list[dict[str, object]], ) -> str: conversation_name = ( normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}" ) first_activity, last_activity = conversation_preview_bounds(documents) entries = [ entry for document in documents for entry in chat_preview_entries_for_document(document) ] body_text = normalize_whitespace( "\n\n".join( str(document.get("text_content") or "") for document in documents if normalize_whitespace(str(document.get("text_content") or "")) ) ) headers = { "Author": "", "Participants": conversation_preview_participants(documents) or "", "Started": format_chat_preview_timestamp(first_activity) or first_activity or "", "Updated": format_chat_preview_timestamp(last_activity) or last_activity or "", "Title": conversation_name, "Messages": str(len(entries) or len(documents)), } return build_chat_preview_html( headers, body_text, document_title=conversation_name, entries=entries, ) def chat_message_preview_rows( paths: dict[str, Path], *, document: dict[str, object], preserved_preview_rows: list[dict[str, object]], entry_rel_path: str, created_at: str, ) -> list[dict[str, object]]: preview_html = build_chat_document_preview_html(document) preferred_rows = sorted( preserved_preview_rows, key=lambda row: ( 0 if normalize_whitespace(str(row.get("preview_type") or "")).lower() == "html" else 1, int(row.get("ordinal", 0)), ), ) preferred_html_rows = [ row for row in preferred_rows if normalize_whitespace(str(row.get("preview_type") or "")).lower() == "html" ] if preferred_html_rows: target_row = dict( next( ( row for row in preferred_html_rows if normalize_whitespace(str(row.get("label") or "")).lower() == "message" ), preferred_html_rows[0], ) ) target_rel_path = normalize_whitespace(str(target_row.get("rel_preview_path") or "")) or entry_rel_path entry_abs_path = paths["state_dir"] / target_rel_path entry_abs_path.parent.mkdir(parents=True, exist_ok=True) entry_abs_path.write_text( preview_html, encoding="utf-8", ) target_row["rel_preview_path"] = target_rel_path target_row["target_fragment"] = None target_row["label"] = "message" target_row["ordinal"] = 0 target_row["created_at"] = created_at return [target_row] entry_abs_path = paths["state_dir"] / entry_rel_path entry_abs_path.parent.mkdir(parents=True, exist_ok=True) entry_abs_path.write_text( preview_html, encoding="utf-8", ) return [ { "rel_preview_path": entry_rel_path, "preview_type": "html", "target_fragment": None, "label": "message", "ordinal": 0, "created_at": created_at, } ] def render_conversation_chat_body_html(text_content: str) -> str: chat_entries = iter_chat_transcript_entries(text_content, max_lines=4000) if not chat_entries: return f"
{html.escape(text_content)}
" rendered_entries: list[str] = [] for entry in chat_entries: speaker = normalize_whitespace(str(entry.get("speaker") or "")) or "Unknown" body = normalize_whitespace(str(entry.get("body") or "")) if not body: continue timestamp_label = format_chat_preview_timestamp(entry.get("timestamp")) or "" timestamp_html = f'[{html.escape(timestamp_label)}]' if timestamp_label else "" avatar_label = chat_avatar_initials(speaker) avatar_background, avatar_foreground = chat_avatar_colors(speaker) avatar_html = build_chat_avatar_svg(avatar_label, avatar_background, avatar_foreground, speaker) rendered_entries.append( "
" f"{avatar_html}" "
" "
" f"{html.escape(speaker)}" f"{timestamp_html}" "
" f"
{html.escape(body)}
" "
" "
" ) if not rendered_entries: return f"
{html.escape(text_content)}
" return ( "
" f"{''.join(rendered_entries)}" "
" "
" "Full transcript" f"
{html.escape(text_content)}
" "
" ) def render_conversation_document_section( document: dict[str, object], *, current_segment_href: str, doc_target_hrefs: dict[int, str], attachment_links_by_document_id: dict[int, list[dict[str, str]]] | None = None, target_preview_rel_path: object = None, ) -> str: document_id = int(document["id"]) anchor = conversation_preview_anchor(document_id) heading = conversation_preview_document_heading(document) kind_label = conversation_preview_document_kind(document) timestamp_label = format_chat_preview_timestamp(conversation_preview_primary_timestamp(document)) or "" text_content = str(document.get("text_content") or "") standalone_preview_body_html = document.get("standalone_preview_body_html") content_type = normalize_whitespace(str(document.get("content_type") or "")) metadata_pairs: list[tuple[str, object]] = [ ("Control number", document.get("control_number")), ("Created", timestamp_label), ] if content_type == "Email": metadata_pairs.extend( [ ("Author", document.get("author")), ("Recipients", document.get("recipients")), ] ) else: metadata_pairs.extend( [ ("Participants", document.get("participants")), ("From", document.get("author")), ("To", document.get("recipients")), ] ) metadata_items: list[str] = [] for label, value in metadata_pairs: normalized = normalize_whitespace(str(value or "")) if not normalized: continue metadata_items.append( f"
{html.escape(label)}
{html.escape(normalized)}
" ) parent_document_id = document.get("parent_document_id") if parent_document_id is not None: parent_id = int(parent_document_id) parent_href = doc_target_hrefs.get(parent_id) parent_label = normalize_whitespace( str(document.get("parent_control_number") or document.get("parent_title") or f"Document {parent_id}") ) if parent_href and parent_label: metadata_items.append( "" ) body_html = ( render_conversation_chat_body_html(text_content) if normalize_whitespace(str(document.get("content_type") or "")) == "Chat" else preview_body_html_for_output( document, str(standalone_preview_body_html), target_preview_rel_path=target_preview_rel_path, ) if isinstance(standalone_preview_body_html, str) and standalone_preview_body_html.strip() else f"
{html.escape(text_content or 'No extracted text available.')}
" ) attachment_links_html = "" if attachment_links_by_document_id: attachment_links = attachment_links_by_document_id.get(document_id) or [] attachment_links_html = render_html_preview_attachment_links(attachment_links) metadata_html = ( f'
{"".join(metadata_items)}
' if metadata_items else "" ) # Permalinks are intentionally omitted: Cowork's preview iframe blocks link # clicks, so a fragment anchor here wouldn't actually jump anywhere useful. _ = current_segment_href return "".join( [ f'
', '
', "
", f'
{html.escape(kind_label)}
', f"

{html.escape(heading)}

", "
", "
", metadata_html, f'
{body_html}
', attachment_links_html, "
", ] ) def build_conversation_preview_head_html() -> str: return ( "" ) def conversation_segment_position_label(segment_index: int, segment_count: int) -> str | None: if segment_count <= 1: return None return f"Segment {segment_index + 1} of {segment_count}" def build_conversation_segment_banner_html( *, segment_label: str, segment_index: int, segment_count: int, document_count: int, ) -> str: position_label = conversation_segment_position_label(segment_index, segment_count) if position_label is None: return "" document_label = f"{document_count} document{'s' if document_count != 1 else ''} in this segment" return ( '
' f'

{html.escape(position_label)}

' f'

{html.escape(segment_label)}

' f'

{html.escape(document_label)}

' "
" ) def build_conversation_toc_html( conversation_row: sqlite3.Row, *, documents: list[dict[str, object]], segment_items: list[dict[str, object]], ) -> str: headers = { "Conversation": normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}", "Type": normalize_whitespace(str(conversation_row["conversation_type"] or "")), "Documents": str(len(documents)), "Segments": str(len(segment_items)), } cards: list[str] = [] for segment in segment_items: doc_entries = "".join( f"
  • {html.escape(conversation_preview_document_heading(document))}
    • " for document in segment["documents"] ) cards.append( "
    " f"

    {html.escape(str(segment['label']))}

    " f"

    {len(segment['documents'])} document{'s' if len(segment['documents']) != 1 else ''}

    " f"{'
      ' + doc_entries + '
    ' if doc_entries else ''}" "
    " ) return build_html_preview( headers, body_html=f"
    {''.join(cards)}
    ", document_title=headers["Conversation"] or "Conversation", head_html=build_conversation_preview_head_html(), heading="Conversation Contents", ) def build_conversation_segment_html( conversation_row: sqlite3.Row, *, segment_label: str, segment_index: int, segment_count: int, segment_items: list[dict[str, object]], current_segment_rel_path: str, doc_target_hrefs: dict[int, str], attachment_links_by_document_id: dict[int, list[dict[str, str]]] | None = None, ) -> str: current_file_name = Path(current_segment_rel_path).name segment_documents = segment_items[segment_index]["documents"] if normalize_whitespace(str(conversation_row["conversation_type"] or "")).lower() == "email": thread_title = ( normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}" ) return build_email_thread_preview_html( thread_title=thread_title, documents=segment_documents, page_title=f"{thread_title} - {segment_label}", segment_label=( f"{segment_label} ({conversation_segment_position_label(segment_index, segment_count)})" if conversation_segment_position_label(segment_index, segment_count) else segment_label ), segment_count=segment_count, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=current_segment_rel_path, strip_quoted_history=True, ) headers = { "Conversation": normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}", "Type": normalize_whitespace(str(conversation_row["conversation_type"] or "")), "Segment": segment_label, "Documents": str(len(segment_documents)), } sections = "".join( render_conversation_document_section( document, current_segment_href=current_file_name, doc_target_hrefs=doc_target_hrefs, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=current_segment_rel_path, ) for document in segment_documents ) segment_banner_html = build_conversation_segment_banner_html( segment_label=segment_label, segment_index=segment_index, segment_count=segment_count, document_count=len(segment_documents), ) return build_html_preview( headers, body_html=f"
    {segment_banner_html}{sections}
    ", document_title=f"{headers['Conversation']} - {segment_label}", head_html=build_conversation_preview_head_html(), heading=segment_label, ) def build_conversation_entry_html( conversation_row: sqlite3.Row, *, document: dict[str, object], document_heading: str, segment_label: str | None = None, attachment_links: list[dict[str, str]] | None = None, position_index: int | None = None, total_count: int | None = None, target_preview_rel_path: object = None, ) -> str: conversation_name = ( normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}" ) document_id = int(document["id"]) section_html = render_conversation_document_section( document, current_segment_href="entry", doc_target_hrefs={document_id: f"#{conversation_preview_anchor(document_id)}"}, attachment_links_by_document_id=( {document_id: attachment_links} if attachment_links else None ), target_preview_rel_path=target_preview_rel_path, ) headers: dict[str, str] = { "Conversation": conversation_name, "Type": normalize_whitespace(str(conversation_row["conversation_type"] or "")), } if segment_label: headers["Segment"] = segment_label if position_index is not None and total_count is not None: headers["Document"] = f"{position_index} of {total_count}" return build_html_preview( headers, body_html=f"
    {section_html}
    ", document_title=document_heading or conversation_name or "Conversation document", head_html=build_conversation_preview_head_html(), heading=document_heading or "Conversation document", ) def build_email_conversation_full_html( conversation_row: sqlite3.Row, *, documents: list[dict[str, object]], segment_items: list[dict[str, object]], attachment_links_by_document_id: dict[int, list[dict[str, str]]] | None = None, target_preview_rel_path: object = None, ) -> str: thread_title = ( normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}" ) if len(segment_items) <= 1: return build_email_thread_preview_html( thread_title=thread_title, documents=documents, page_title=thread_title, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=target_preview_rel_path, strip_quoted_history=True, ) segment_sections: list[str] = [] for index, segment in enumerate(segment_items): segment_documents = list(segment["documents"]) message_cards = "".join( build_email_message_card_html( document, attachment_links=(attachment_links_by_document_id or {}).get(int(document["id"]), []), target_preview_rel_path=target_preview_rel_path, strip_quoted_history=True, ) for document in segment_documents ) segment_sections.append( '
    ' '
    ' f'

    {html.escape(conversation_segment_position_label(index, len(segment_items)) or "")}

    ' f'

    {html.escape(str(segment["label"]))}

    ' f"{build_email_thread_summary_html(segment_documents)}" "
    " f'
    {message_cards}
    ' "
    " ) return ( "" "" '' f"{html.escape(thread_title)}" f"{build_email_preview_head_html()}" "" "" '
    ' '
    ' f'

    {html.escape(thread_title)}

    ' f"{build_email_thread_summary_html(documents)}" "
    " f'
    {"".join(segment_sections)}
    ' "
    " "" ) def build_conversation_full_html( conversation_row: sqlite3.Row, *, documents: list[dict[str, object]], segment_items: list[dict[str, object]], attachment_links_by_document_id: dict[int, list[dict[str, str]]] | None = None, ) -> str: full_rel_path = conversation_preview_full_rel_path(int(conversation_row["id"])) if normalize_whitespace(str(conversation_row["conversation_type"] or "")).lower() == "email": return build_email_conversation_full_html( conversation_row, documents=documents, segment_items=segment_items, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=full_rel_path, ) if conversation_documents_are_chat(documents): return build_chat_conversation_preview_html( conversation_row, documents=documents, ) conversation_name = ( normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {int(conversation_row['id'])}" ) headers = { "Conversation": conversation_name, "Type": normalize_whitespace(str(conversation_row["conversation_type"] or "")), "Documents": str(len(documents)), "Segments": str(len(segment_items)), } doc_target_hrefs = { int(document["id"]): f"#{conversation_preview_anchor(int(document['id']))}" for document in documents } segment_sections = "".join( ( '
    ' f"{build_conversation_segment_banner_html(segment_label=str(segment['label']), segment_index=index, segment_count=len(segment_items), document_count=len(segment['documents']))}" + "".join( render_conversation_document_section( document, current_segment_href="conversation", doc_target_hrefs=doc_target_hrefs, attachment_links_by_document_id=attachment_links_by_document_id, target_preview_rel_path=full_rel_path, ) for document in segment["documents"] ) + "
    " ) for index, segment in enumerate(segment_items) ) return build_html_preview( headers, body_html=f'
    {segment_sections}
    ', document_title=conversation_name, head_html=build_conversation_preview_head_html(), heading=conversation_name, ) def conversation_document_uses_entry_preview(document: dict[str, object]) -> bool: return normalize_whitespace(str(document.get("content_type") or "")).lower() == "chat" def document_content_type_is_chat(content_type: object) -> bool: return normalize_whitespace(str(content_type or "")).lower() == "chat" def document_content_type_is_email(content_type: object) -> bool: return normalize_whitespace(str(content_type or "")).lower() == "email" def text_revision_pointers_differ( source_text_revision_id: object, active_search_text_revision_id: object, ) -> bool: if source_text_revision_id is None or active_search_text_revision_id is None: return False return int(active_search_text_revision_id) != int(source_text_revision_id) def document_active_text_differs_from_source(document: dict[str, object]) -> bool: return text_revision_pointers_differ( document.get("source_text_revision_id"), document.get("active_search_text_revision_id"), ) def document_active_text_source_kind(document: dict[str, object]) -> str: return normalize_whitespace(str(document.get("active_text_source_kind") or "")).lower() def document_active_text_should_render_above_preserved_email_body( document: dict[str, object], ) -> bool: return ( document_active_text_differs_from_source(document) and document_active_text_source_kind(document) == "ocr" ) def load_document_preview_text( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, storage_rel_path: object, prefer_search_chunks: bool = False, ) -> str: chunk_rows: list[sqlite3.Row] | None = None if prefer_search_chunks: chunk_rows = connection.execute( """ SELECT text_content FROM document_chunks WHERE document_id = ? ORDER BY chunk_index ASC """, (document_id,), ).fetchall() if chunk_rows: return "\n".join(str(row["text_content"] or "") for row in chunk_rows) text_content = read_text_revision_body(paths, str(storage_rel_path or "") or None) if text_content is not None: return text_content if chunk_rows is None: chunk_rows = connection.execute( """ SELECT text_content FROM document_chunks WHERE document_id = ? ORDER BY chunk_index ASC """, (document_id,), ).fetchall() return "\n".join(str(row["text_content"] or "") for row in chunk_rows) def load_preview_documents( connection: sqlite3.Connection, paths: dict[str, Path], *, document_ids: list[int] | None = None, conversation_id: int | None = None, include_attachment_children: bool = False, require_dataset_membership: bool = False, ) -> list[dict[str, object]]: if document_ids is not None and conversation_id is not None: raise RetrieverError("load_preview_documents accepts either document_ids or conversation_id, not both.") if document_ids is None and conversation_id is None: raise RetrieverError("load_preview_documents requires document_ids or conversation_id.") where_clauses = ["d.lifecycle_status NOT IN ('missing', 'deleted')"] parameters: list[object] = [] if document_ids is not None: normalized_document_ids = [int(document_id) for document_id in document_ids] if not normalized_document_ids: return [] where_clauses.append(f"d.id IN ({', '.join('?' for _ in normalized_document_ids)})") parameters.extend(normalized_document_ids) else: where_clauses.append("d.conversation_id = ?") parameters.append(int(conversation_id)) if not include_attachment_children: where_clauses.append("COALESCE(d.child_document_kind, '') != ?") parameters.append(CHILD_DOCUMENT_KIND_ATTACHMENT) if require_dataset_membership: where_clauses.append( "EXISTS (SELECT 1 FROM dataset_documents dd WHERE dd.document_id = d.id)" ) rows = connection.execute( f""" SELECT d.id, d.rel_path, d.control_number, d.parent_document_id, d.child_document_kind, d.file_name, d.file_type, d.content_type, d.date_created, d.date_modified, d.title, d.subject, d.author, d.participants, d.recipients, d.source_kind, d.source_rel_path, d.source_item_id, d.source_folder_path, d.root_message_key, d.conversation_id, d.source_text_revision_id, d.active_search_text_revision_id, d.active_text_source_kind, parent.control_number AS parent_control_number, parent.title AS parent_title, source_tr.storage_rel_path AS source_text_storage_rel_path, active_tr.storage_rel_path AS active_text_storage_rel_path FROM documents d LEFT JOIN documents parent ON parent.id = d.parent_document_id LEFT JOIN text_revisions source_tr ON source_tr.id = d.source_text_revision_id LEFT JOIN text_revisions active_tr ON active_tr.id = COALESCE(d.active_search_text_revision_id, d.source_text_revision_id) WHERE {' AND '.join(where_clauses)} ORDER BY d.id ASC """, tuple(parameters), ).fetchall() documents: list[dict[str, object]] = [] for row in rows: documents.append( { key: row[key] for key in row.keys() } | { "text_content": load_document_preview_text( connection, paths, document_id=int(row["id"]), storage_rel_path=row["active_text_storage_rel_path"], prefer_search_chunks=document_content_type_is_chat(row["content_type"]), ) } ) preserved_preview_rows_by_document_id = load_preserved_preview_rows_by_document_id( connection, paths, [int(document["id"]) for document in documents], ) for document in documents: if normalize_whitespace(str(document.get("content_type") or "")).lower() == "chat": continue body_payload = load_document_preview_body_payload( paths, preserved_preview_rows_by_document_id.get(int(document["id"]), []), ) if body_payload is None and document_content_type_is_email(document.get("content_type")): body_payload = load_source_backed_email_preview_body_payload(paths, document) if body_payload: document["standalone_preview_body_html"] = body_payload["body_html"] document["standalone_preview_rel_path"] = body_payload["rel_preview_path"] documents.sort(key=conversation_preview_sort_key) return documents def refresh_conversation_previews( connection: sqlite3.Connection, paths: dict[str, Path], conversation_ids: list[int] | None = None, ) -> int: target_conversation_ids = ( sorted(dict.fromkeys(int(conversation_id) for conversation_id in conversation_ids)) if conversation_ids is not None else list_active_conversation_ids(connection) ) refreshed = 0 for conversation_id in target_conversation_ids: conversation_row = connection.execute( "SELECT * FROM conversations WHERE id = ?", (conversation_id,), ).fetchone() if conversation_row is None: continue documents = load_preview_documents(connection, paths, conversation_id=conversation_id) if not documents: continue segment_mode = conversation_preview_segment_mode(conversation_row["conversation_type"], documents) writes_aggregate_artifacts = conversation_preview_writes_aggregate_artifacts( conversation_row["conversation_type"], documents, segment_mode=segment_mode, ) segment_documents: dict[str, list[dict[str, object]]] = defaultdict(list) for document in documents: segment_documents[conversation_preview_segment_key(document, segment_mode=segment_mode)].append(document) segment_items = [ { "segment_key": segment_key, "label": conversation_preview_segment_label(segment_key, segment_mode=segment_mode), "segment_rel_path": conversation_preview_segment_rel_path(conversation_id, segment_key), "documents": sorted(items, key=conversation_preview_sort_key), } for segment_key, items in sorted(segment_documents.items(), key=lambda item: item[0]) ] document_ids = [int(document["id"]) for document in documents] preserved_preview_rows_by_document_id = load_preserved_preview_rows_by_document_id( connection, paths, document_ids, ) writes_segment_artifacts = ( writes_aggregate_artifacts and len(segment_items) > 1 and not conversation_documents_are_chat(documents) ) toc_rel_path = conversation_preview_toc_rel_path(conversation_id) full_rel_path = conversation_preview_full_rel_path(conversation_id) full_preview_existed_before = (paths["state_dir"] / full_rel_path).exists() def aggregate_preview_rel_path(segment: dict[str, object]) -> str: return str(segment["segment_rel_path"]) if writes_segment_artifacts else full_rel_path doc_target_hrefs = { int(document["id"]): f"{Path(aggregate_preview_rel_path(segment)).name}#{conversation_preview_anchor(int(document['id']))}" for segment in segment_items for document in segment["documents"] } email_message_position_by_document_id = { int(document["id"]): index + 1 for index, document in enumerate( [document for document in documents if normalize_whitespace(str(document.get("content_type") or "")).lower() == "email"] ) } if writes_aggregate_artifacts: full_abs_path = paths["state_dir"] / full_rel_path full_abs_path.parent.mkdir(parents=True, exist_ok=True) full_attachment_links = conversation_attachment_links_by_document_id( connection, paths, segment_preview_path=full_abs_path, documents=documents, ) full_abs_path.write_text( build_conversation_full_html( conversation_row, documents=documents, segment_items=segment_items, attachment_links_by_document_id=full_attachment_links, ), encoding="utf-8", ) if writes_segment_artifacts: toc_abs_path = paths["state_dir"] / toc_rel_path toc_abs_path.parent.mkdir(parents=True, exist_ok=True) toc_abs_path.write_text( build_conversation_toc_html( conversation_row, documents=documents, segment_items=segment_items, ), encoding="utf-8", ) for index, segment in enumerate(segment_items): segment_abs_path = paths["state_dir"] / str(segment["segment_rel_path"]) segment_abs_path.parent.mkdir(parents=True, exist_ok=True) attachment_links = conversation_attachment_links_by_document_id( connection, paths, segment_preview_path=segment_abs_path, documents=list(segment["documents"]), ) segment_abs_path.write_text( build_conversation_segment_html( conversation_row, segment_label=str(segment["label"]), segment_index=index, segment_count=len(segment_items), segment_items=segment_items, current_segment_rel_path=str(segment["segment_rel_path"]), doc_target_hrefs=doc_target_hrefs, attachment_links_by_document_id=attachment_links, ), encoding="utf-8", ) created_at = utc_now() document_ids = [int(document["id"]) for document in documents] previous_preview_paths = [ str(row["rel_preview_path"]) for row in connection.execute( f""" SELECT DISTINCT rel_preview_path FROM document_previews WHERE document_id IN ({", ".join("?" for _ in document_ids)}) """, document_ids, ).fetchall() ] if full_preview_existed_before: previous_preview_paths.append(full_rel_path) for segment in segment_items: segment_rel_path = aggregate_preview_rel_path(segment) segment_abs_path = paths["state_dir"] / segment_rel_path entry_attachment_links = conversation_attachment_links_by_document_id( connection, paths, segment_preview_path=segment_abs_path, documents=list(segment["documents"]), ) for document in segment["documents"]: document_id = int(document["id"]) preserved_preview_rows = preserved_preview_rows_by_document_id.get(document_id, []) if normalize_whitespace(str(document.get("content_type") or "")).lower() == "email": rewrite_preserved_email_message_preview( paths, document=document, preview_rows=preserved_preview_rows, conversation_row=conversation_row, conversation_documents=documents, position_index=( email_message_position_by_document_id.get(document_id) if len(documents) > 1 else None ), thread_rel_path=( segment_rel_path if writes_aggregate_artifacts else None ), attachment_links=entry_attachment_links.get(document_id) or [], ) preview_rows: list[dict[str, object]] = [] if conversation_document_uses_entry_preview(document): entry_rel_path = conversation_preview_entry_rel_path(conversation_id, document_id) preview_rows.extend( chat_message_preview_rows( paths, document=document, preserved_preview_rows=preserved_preview_rows, entry_rel_path=entry_rel_path, created_at=created_at, ) ) if writes_aggregate_artifacts: preview_rows.append( { "rel_preview_path": full_rel_path, "preview_type": "html", "target_fragment": None, "label": "conversation", "ordinal": len(preview_rows), "created_at": created_at, } ) rebased_preview_rows = [] else: rebased_preview_rows = rebase_preserved_preview_rows( preserved_preview_rows, start_ordinal=0, created_at=created_at, ) preview_rows.extend(rebased_preview_rows) if writes_aggregate_artifacts: preview_rows.append( { "rel_preview_path": segment_rel_path, "preview_type": "html", "target_fragment": None, "label": "segment", "ordinal": len(preview_rows), "created_at": created_at, } ) if writes_segment_artifacts: preview_rows.append( { "rel_preview_path": toc_rel_path, "preview_type": "html", "target_fragment": None, "label": "contents", "ordinal": len(preview_rows), "created_at": created_at, } ) replace_document_preview_rows( connection, int(document["id"]), preview_rows if not conversation_document_uses_entry_preview(document) else [*preview_rows, *rebased_preview_rows], ) cleanup_unreferenced_preview_files(paths, connection, previous_preview_paths) refreshed += 1 return refreshed def mark_missing_documents( connection: sqlite3.Connection, scanned_rel_paths: set[str], scan_scope: dict[str, object] | None = None, ) -> int: occurrence_rows = connection.execute( """ SELECT id, document_id, rel_path, lifecycle_status FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND lifecycle_status != 'deleted' """ , (FILESYSTEM_SOURCE_KIND,)).fetchall() missing_occurrence_ids = [ int(row["id"]) for row in occurrence_rows if ingest_scan_scope_contains_rel_path(scan_scope, row["rel_path"]) and row["rel_path"] not in scanned_rel_paths and row["lifecycle_status"] != "missing" ] if not missing_occurrence_ids: return 0 now = utc_now() placeholders = ", ".join("?" for _ in missing_occurrence_ids) connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'missing', updated_at = ? WHERE lifecycle_status != 'deleted' AND (id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders})) """, [now, *missing_occurrence_ids, *missing_occurrence_ids], ) affected_document_ids = { int(row["document_id"]) for row in connection.execute( f""" SELECT DISTINCT document_id FROM document_occurrences WHERE id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders}) """, [*missing_occurrence_ids, *missing_occurrence_ids], ).fetchall() } for document_id in affected_document_ids: refresh_source_backed_dataset_memberships_for_document(connection, document_id) refresh_document_from_occurrences(connection, document_id) connection.commit() return len(missing_occurrence_ids) def select_attachment_match_candidate( buckets: dict[object, list[sqlite3.Row]], key: object, matched_ids: set[int], ) -> sqlite3.Row | None: candidates = buckets.get(key) or [] while candidates and int(candidates[0]["id"]) in matched_ids: candidates.pop(0) if not candidates: return None candidate = candidates.pop(0) matched_ids.add(int(candidate["id"])) return candidate def match_attachment_rows( existing_rows: list[sqlite3.Row], attachments: list[dict[str, object]], ) -> tuple[list[tuple[dict[str, object], sqlite3.Row | None]], list[sqlite3.Row]]: candidate_rows = [row for row in existing_rows if row["lifecycle_status"] != "deleted"] by_name_hash: dict[tuple[str | None, str], list[sqlite3.Row]] = defaultdict(list) by_hash: dict[str | None, list[sqlite3.Row]] = defaultdict(list) for row in candidate_rows: file_hash = row["file_hash"] file_name = str(row["file_name"] or "").lower() by_name_hash[(file_hash, file_name)].append(row) by_hash[file_hash].append(row) matched_ids: set[int] = set() matches: list[tuple[dict[str, object], sqlite3.Row | None]] = [] for attachment in attachments: file_hash = attachment.get("file_hash") file_name = str(attachment.get("file_name") or "").lower() existing_row = select_attachment_match_candidate(by_name_hash, (file_hash, file_name), matched_ids) if existing_row is None: existing_row = select_attachment_match_candidate(by_hash, file_hash, matched_ids) matches.append((attachment, existing_row)) removed_rows = [row for row in candidate_rows if int(row["id"]) not in matched_ids] return matches, removed_rows def extract_attachment_document(path: Path) -> dict[str, object]: try: return extract_document(path, include_attachments=False) except Exception: return build_fallback_extract(path) def extract_attachment_document_with_overrides( path: Path, attachment: dict[str, object], ) -> dict[str, object]: extracted = extract_attachment_document(path) drive_record = attachment.get("gmail_drive_record") if isinstance(drive_record, dict): extracted = apply_gmail_drive_export_metadata( dict(extracted), drive_record=dict(drive_record), ) return extracted def reconcile_attachment_documents( connection: sqlite3.Connection, paths: dict[str, Path], parent_document_id: int, parent_rel_path: str, control_number_batch: int, control_number_family_sequence: int, attachments: list[dict[str, object]], dataset_memberships: list[tuple[int, int | None]] | None = None, ) -> None: parent_row = connection.execute( "SELECT dataset_id, conversation_id, conversation_assignment_mode FROM documents WHERE id = ?", (parent_document_id,), ).fetchone() parent_occurrence = select_preferred_occurrence(active_occurrence_rows_for_document(connection, parent_document_id)) parent_custodian = normalize_whitespace(str(parent_occurrence["custodian"] or "")) or None if parent_occurrence is not None else None parent_dataset_id = int(parent_row["dataset_id"]) if parent_row is not None and parent_row["dataset_id"] is not None else None parent_conversation_id = int(parent_row["conversation_id"]) if parent_row is not None and parent_row["conversation_id"] is not None else None parent_conversation_assignment_mode = ( str(parent_row["conversation_assignment_mode"]) if parent_row is not None and parent_row["conversation_assignment_mode"] is not None else CONVERSATION_ASSIGNMENT_MODE_AUTO ) existing_rows = connection.execute( """ SELECT * FROM documents WHERE parent_document_id = ? ORDER BY id ASC """, (parent_document_id,), ).fetchall() active_occurrences_by_document_id: dict[int, sqlite3.Row] = {} if existing_rows: existing_document_ids = [int(row["id"]) for row in existing_rows] placeholders = ", ".join("?" for _ in existing_document_ids) occurrence_rows = connection.execute( f""" SELECT * FROM document_occurrences WHERE document_id IN ({placeholders}) AND lifecycle_status != 'deleted' ORDER BY id ASC """, existing_document_ids, ).fetchall() grouped_occurrences: dict[int, list[sqlite3.Row]] = defaultdict(list) for occurrence_row in occurrence_rows: grouped_occurrences[int(occurrence_row["document_id"])].append(occurrence_row) active_occurrences_by_document_id = { document_id: select_preferred_occurrence(rows) or rows[0] for document_id, rows in grouped_occurrences.items() if rows } matches, removed_rows = match_attachment_rows(existing_rows, attachments) next_new_attachment_sequence = next_attachment_sequence(connection, parent_document_id) for attachment, existing_row in matches: if existing_row is None: attachment_sequence = next_new_attachment_sequence next_new_attachment_sequence += 1 control_number = format_control_number(control_number_batch, control_number_family_sequence, attachment_sequence) else: attachment_sequence = int(existing_row["control_number_attachment_sequence"]) control_number = str(existing_row["control_number"]) cleanup_document_artifacts(paths, connection, existing_row) child_rel_path, child_path = write_attachment_blob( paths, parent_rel_path, control_number, str(attachment["file_name"]), bytes(attachment["payload"]), ) extracted = apply_manual_locks( existing_row, extract_attachment_document_with_overrides(child_path, attachment), ) document_id = upsert_document_row( connection, child_rel_path, child_path, existing_row, extracted, existing_occurrence_row=( active_occurrences_by_document_id.get(int(existing_row["id"])) if existing_row is not None else None ), file_name=str(attachment["file_name"]), parent_document_id=parent_document_id, control_number=control_number, dataset_id=parent_dataset_id, conversation_id=parent_conversation_id, conversation_assignment_mode=parent_conversation_assignment_mode, control_number_batch=control_number_batch, control_number_family_sequence=control_number_family_sequence, control_number_attachment_sequence=attachment_sequence, custodian_override=parent_custodian, ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=existing_row, ) preview_rows = write_preview_artifacts(paths, child_rel_path, list(extracted.get("preview_artifacts", []))) chunks = extracted_search_chunks(extracted) replace_document_related_rows( connection, document_id, extracted | {"file_name": str(attachment["file_name"])}, chunks, preview_rows, ) for membership_dataset_id, membership_source_id in dataset_memberships or []: ensure_dataset_document_membership( connection, dataset_id=membership_dataset_id, document_id=document_id, dataset_source_id=membership_source_id, ) for row in removed_rows: cleanup_document_artifacts(paths, connection, row) delete_document_related_rows(connection, int(row["id"])) connection.execute( """ UPDATE documents SET lifecycle_status = 'deleted', updated_at = ? WHERE id = ? """, (utc_now(), row["id"]), ) connection.execute( """ UPDATE document_occurrences SET lifecycle_status = 'deleted', updated_at = ? WHERE document_id = ? """, (utc_now(), row["id"]), ) sync_document_attachment_preview_links(connection, paths, parent_document_id) def get_container_source_row( connection: sqlite3.Connection, source_kind: str, source_rel_path: str, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM container_sources WHERE source_kind = ? AND source_rel_path = ? """, (source_kind, source_rel_path), ).fetchone() def container_source_scan_completed(row: sqlite3.Row | None) -> bool: if row is None: return False started_at = parse_utc_timestamp(row["last_scan_started_at"]) completed_at = parse_utc_timestamp(row["last_scan_completed_at"]) return started_at is not None and completed_at is not None and completed_at >= started_at def write_container_source_scan_started( connection: sqlite3.Connection, *, dataset_id: int | None, source_kind: str, source_rel_path: str, file_size: int | None, file_mtime: str | None, file_hash: str | None, scan_started_at: str, ) -> None: connection.execute( """ INSERT INTO container_sources ( dataset_id, source_kind, source_rel_path, file_size, file_mtime, file_hash, last_scan_started_at ) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(source_rel_path) DO UPDATE SET dataset_id = excluded.dataset_id, source_kind = excluded.source_kind, file_size = excluded.file_size, file_mtime = excluded.file_mtime, file_hash = excluded.file_hash, last_scan_started_at = excluded.last_scan_started_at """, (dataset_id, source_kind, source_rel_path, file_size, file_mtime, file_hash, scan_started_at), ) def write_container_source_scan_completed( connection: sqlite3.Connection, *, dataset_id: int | None, source_kind: str, source_rel_path: str, file_size: int | None, file_mtime: str | None, file_hash: str | None, message_count: int, scan_started_at: str, scan_completed_at: str, ) -> None: connection.execute( """ INSERT INTO container_sources ( dataset_id, source_kind, source_rel_path, file_size, file_mtime, file_hash, message_count, last_scan_started_at, last_scan_completed_at, last_ingested_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(source_rel_path) DO UPDATE SET dataset_id = excluded.dataset_id, source_kind = excluded.source_kind, file_size = excluded.file_size, file_mtime = excluded.file_mtime, file_hash = excluded.file_hash, message_count = excluded.message_count, last_scan_started_at = excluded.last_scan_started_at, last_scan_completed_at = excluded.last_scan_completed_at, last_ingested_at = excluded.last_ingested_at """, ( dataset_id, source_kind, source_rel_path, file_size, file_mtime, file_hash, message_count, scan_started_at, scan_completed_at, scan_completed_at, ), ) def mark_container_source_documents_active( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, seen_at: str, ) -> int: root_occurrence_ids = [ int(row["id"]) for row in container_root_occurrence_rows_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ] if not root_occurrence_ids: return 0 placeholders = ", ".join("?" for _ in root_occurrence_ids) cursor = connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'active', last_seen_at = ?, updated_at = ? WHERE lifecycle_status != 'deleted' AND (id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders})) """, [seen_at, seen_at, *root_occurrence_ids, *root_occurrence_ids], ) affected_document_ids = container_document_ids_for_root_occurrence_ids(connection, root_occurrence_ids) for document_id in sorted(affected_document_ids): refresh_source_backed_dataset_memberships_for_document(connection, document_id) refresh_document_from_occurrences(connection, document_id) return int(cursor.rowcount or 0) def assign_dataset_to_container_documents( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, dataset_id: int, dataset_source_id: int | None = None, ) -> None: _ = (dataset_id, dataset_source_id) for document_id in sorted( container_document_ids_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ): refresh_source_backed_dataset_memberships_for_document(connection, document_id) def existing_container_entries_by_source_item( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, ) -> dict[str, dict[str, sqlite3.Row]]: occurrence_rows = connection.execute( """ SELECT * FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND source_rel_path = ? AND source_item_id IS NOT NULL AND lifecycle_status != 'deleted' ORDER BY id ASC """, (source_kind, source_rel_path), ).fetchall() if not occurrence_rows: return {} document_ids = sorted({int(row["document_id"]) for row in occurrence_rows}) placeholders = ", ".join("?" for _ in document_ids) document_rows = connection.execute( f""" SELECT * FROM documents WHERE id IN ({placeholders}) AND parent_document_id IS NULL ORDER BY id ASC """, document_ids, ).fetchall() documents_by_id = {int(row["id"]): row for row in document_rows} rows_by_source_item: dict[str, list[sqlite3.Row]] = defaultdict(list) for occurrence_row in occurrence_rows: document_id = int(occurrence_row["document_id"]) if document_id not in documents_by_id: continue rows_by_source_item[str(occurrence_row["source_item_id"])].append(occurrence_row) entries: dict[str, dict[str, sqlite3.Row]] = {} for source_item_id, source_rows in rows_by_source_item.items(): preferred_row = select_preferred_occurrence(source_rows) or source_rows[0] entries[source_item_id] = { "document_row": documents_by_id[int(preferred_row["document_id"])], "occurrence_row": preferred_row, } return entries def delete_documents_with_only_deleted_occurrences( connection: sqlite3.Connection, paths: dict[str, Path], document_ids: set[int], *, deleted_at: str, ) -> int: deleted = 0 for document_id in sorted(int(value) for value in document_ids): remaining_row = connection.execute( """ SELECT 1 FROM document_occurrences WHERE document_id = ? AND lifecycle_status != 'deleted' LIMIT 1 """, (document_id,), ).fetchone() if remaining_row is not None: continue document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (document_id,), ).fetchone() if document_row is None or document_row["lifecycle_status"] == "deleted": continue cleanup_document_artifacts(paths, connection, document_row) delete_document_related_rows(connection, document_id) connection.execute( """ UPDATE documents SET lifecycle_status = 'deleted', updated_at = ? WHERE id = ? """, (deleted_at, document_id), ) deleted += 1 return deleted def retire_unseen_container_messages( connection: sqlite3.Connection, paths: dict[str, Path], *, source_kind: str, source_rel_path: str, scan_started_at: str, ) -> int: root_occurrence_rows = connection.execute( """ SELECT * FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND source_rel_path = ? AND lifecycle_status != 'deleted' AND (last_seen_at IS NULL OR last_seen_at != ?) ORDER BY id ASC """, (source_kind, source_rel_path, scan_started_at), ).fetchall() if not root_occurrence_rows: return 0 now = utc_now() root_occurrence_ids = [int(row["id"]) for row in root_occurrence_rows] placeholders = ", ".join("?" for _ in root_occurrence_ids) connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'deleted', updated_at = ? WHERE lifecycle_status != 'deleted' AND (id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders})) """, [now, *root_occurrence_ids, *root_occurrence_ids], ) affected_document_ids = container_document_ids_for_root_occurrence_ids(connection, root_occurrence_ids) for document_id in sorted(affected_document_ids): refresh_source_backed_dataset_memberships_for_document(connection, document_id) refresh_document_from_occurrences(connection, document_id) delete_documents_with_only_deleted_occurrences( connection, paths, affected_document_ids, deleted_at=now, ) return len(root_occurrence_rows) def mark_missing_container_documents( connection: sqlite3.Connection, *, source_kind: str, scanned_source_rel_paths: set[str], scan_scope: dict[str, object] | None = None, ) -> tuple[int, int]: source_rows = connection.execute( """ SELECT source_rel_path FROM container_sources WHERE source_kind = ? ORDER BY source_rel_path ASC """, (source_kind,), ).fetchall() sources_missing = 0 documents_missing = 0 now = utc_now() for source_row in source_rows: source_rel_path = str(source_row["source_rel_path"]) if not ingest_scan_scope_contains_rel_path(scan_scope, source_rel_path): continue if source_rel_path in scanned_source_rel_paths: continue root_occurrence_rows = connection.execute( """ SELECT * FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND source_rel_path = ? AND lifecycle_status NOT IN ('missing', 'deleted') ORDER BY id ASC """, (source_kind, source_rel_path), ).fetchall() if not root_occurrence_rows: continue root_occurrence_ids = [int(row["id"]) for row in root_occurrence_rows] placeholders = ", ".join("?" for _ in root_occurrence_ids) connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'missing', updated_at = ? WHERE lifecycle_status NOT IN ('missing', 'deleted') AND (id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders})) """, [now, *root_occurrence_ids, *root_occurrence_ids], ) affected_document_ids = container_document_ids_for_root_occurrence_ids(connection, root_occurrence_ids) for document_id in sorted(affected_document_ids): refresh_source_backed_dataset_memberships_for_document(connection, document_id) refresh_document_from_occurrences(connection, document_id) sources_missing += 1 documents_missing += sum( 1 for document_id in affected_document_ids if connection.execute( "SELECT lifecycle_status FROM documents WHERE id = ?", (document_id,), ).fetchone()["lifecycle_status"] == "missing" ) if sources_missing or documents_missing: connection.commit() return sources_missing, documents_missing def prepare_container_message_item( source_rel_path: str, raw_message: dict[str, object], normalize_message, ) -> dict[str, object]: prepare_started = time.perf_counter() normalized = normalize_message(source_rel_path, raw_message) if normalized is None: return { "skip": True, "source_item_id": normalize_source_item_id(raw_message.get("source_item_id")), "prepare_ms": (time.perf_counter() - prepare_started) * 1000.0, "prepare_chunk_ms": 0.0, } extracted_payload = dict(normalized["extracted"]) attachments = list(extracted_payload.get("attachments", [])) extracted_payload.pop("attachments", None) chunk_started = time.perf_counter() prepared_chunks = extracted_search_chunks(extracted_payload) return { "skip": False, "rel_path": str(normalized["rel_path"]), "file_name": str(normalized["file_name"]), "file_hash": normalized.get("file_hash"), "source_item_id": str(normalized["source_item_id"]), "source_folder_path": ( str(normalized["source_folder_path"]) if normalized.get("source_folder_path") is not None else None ), "extracted_payload": extracted_payload, "attachments": attachments, "prepared_chunks": prepared_chunks, "prepare_ms": (time.perf_counter() - prepare_started) * 1000.0, "prepare_chunk_ms": (time.perf_counter() - chunk_started) * 1000.0, } def iter_prepared_container_message_items( *, source_kind: str, source_rel_path: str, raw_messages: Iterator[dict[str, object]], normalize_message, staging_root: Path | None = None, ) -> Iterator[tuple[dict[str, object], float]]: effective_staging_root = staging_root if staging_root is not None: effective_staging_root = ( Path(staging_root) / "container" / sanitize_storage_filename(source_kind) / sanitize_storage_filename(source_rel_path) ) yield from iter_staged_prepared_items( raw_messages, prepare_item=lambda raw_message: prepare_container_message_item( source_rel_path, raw_message, normalize_message, ), config_benchmark_name="ingest_container_prepare_config", queue_done_benchmark_name="ingest_container_prepare_queue_done", spill_subdir_name="prepared-container", staging_root=effective_staging_root, prepare_workers=ingest_container_prepare_worker_count(), ) def commit_prepared_container_message_in_transaction( connection: sqlite3.Connection, paths: dict[str, Path], prepared_item: dict[str, object], existing_row: sqlite3.Row | None, existing_occurrence_row: sqlite3.Row | None, *, current_ingestion_batch: int | None, dataset_id: int, dataset_source_id: int | None, source_kind: str, source_rel_path: str, file_type_override: str, scan_started_at: str, before_transaction_commit=None, ) -> dict[str, object]: file_hash = ( str(prepared_item["file_hash"]) if prepared_item.get("file_hash") is not None else None ) source_folder_path = prepared_item.get("source_folder_path") extracted_payload = dict(prepared_item["extracted_payload"] or {}) resolved_custodian = normalize_whitespace(str(extracted_payload.get("custodian") or "")) or None if resolved_custodian is None: resolved_custodian = infer_source_custodian( source_kind=source_kind, source_rel_path=source_rel_path, ) if existing_occurrence_row is None and file_hash: exact_duplicate_document = get_document_by_dedupe_key( connection, basis="file_hash", key_value=file_hash, ) if exact_duplicate_document is not None: duplicate_occurrence_id = attach_occurrence_to_existing_document( connection, exact_duplicate_document, existing_occurrence_row=find_active_occurrence_by_source_identity( connection, source_kind=source_kind, custodian=resolved_custodian, source_rel_path=source_rel_path, source_item_id=str(prepared_item["source_item_id"]), ), rel_path=str(prepared_item["rel_path"]), file_name=str(prepared_item["file_name"]), file_type=file_type_override, file_size=None, file_hash=file_hash, source_kind=source_kind, source_rel_path=source_rel_path, source_item_id=str(prepared_item["source_item_id"]), source_folder_path=source_folder_path, custodian=resolved_custodian, occurrence_control_number=str(exact_duplicate_document["control_number"] or ""), ingested_at=scan_started_at, last_seen_at=scan_started_at, updated_at=scan_started_at, ) replace_document_email_threading_row( connection, document_id=int(exact_duplicate_document["id"]), email_threading=extracted_payload.get("email_threading"), ) replace_document_chat_threading_row( connection, document_id=int(exact_duplicate_document["id"]), chat_threading=extracted_payload.get("chat_threading"), ) clone_duplicate_family_child_occurrences( connection, paths, parent_document_id=int(exact_duplicate_document["id"]), parent_occurrence_id=duplicate_occurrence_id, parent_rel_path=str(prepared_item["rel_path"]), custodian=resolved_custodian, ingested_at=scan_started_at, last_seen_at=scan_started_at, updated_at=scan_started_at, ) ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=int(exact_duplicate_document["id"]), dataset_source_id=dataset_source_id, ) duplicate_result = { "action": "new", "current_ingestion_batch": current_ingestion_batch, "document_id": int(exact_duplicate_document["id"]), } if before_transaction_commit is not None: before_transaction_commit(connection, duplicate_result) return duplicate_result existing_document_row = existing_row reused_existing_occurrence_row = existing_occurrence_row superseded_document_id: int | None = None if existing_document_row is not None and file_hash: exact_duplicate_document = get_document_by_dedupe_key( connection, basis="file_hash", key_value=file_hash, ) if ( exact_duplicate_document is not None and int(exact_duplicate_document["id"]) != int(existing_document_row["id"]) and existing_document_row["parent_document_id"] is None and exact_duplicate_document["parent_document_id"] is None ): merge_candidate = evaluate_reconcile_candidate_group( connection, [existing_document_row, exact_duplicate_document], ) if merge_candidate["status"] == "ready": merge_result = apply_evaluated_document_merge_group( connection, paths=paths, merge_basis="ingest:file_hash", merge_group=merge_candidate, ) existing_document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (int(merge_result["survivor_document_id"]),), ).fetchone() if existing_document_row is None: raise RetrieverError( f"Missing survivor document after exact-duplicate merge: " f"{merge_result['survivor_document_id']}" ) if reused_existing_occurrence_row is not None: reused_existing_occurrence_row = connection.execute( "SELECT * FROM document_occurrences WHERE id = ?", (int(reused_existing_occurrence_row["id"]),), ).fetchone() if reused_existing_occurrence_row is None: raise RetrieverError( "Missing source occurrence after exact-duplicate merge during container ingest." ) if reused_existing_occurrence_row is not None: if existing_document_row is None: existing_document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (reused_existing_occurrence_row["document_id"],), ).fetchone() if existing_document_row is None: raise RetrieverError( f"Occurrence {reused_existing_occurrence_row['id']} points at a missing document." ) active_occurrence_rows = active_occurrence_rows_for_document(connection, int(existing_document_row["id"])) if ( len(active_occurrence_rows) > 1 and reused_existing_occurrence_row["file_hash"] != file_hash ): connection.execute( """ UPDATE document_occurrences SET lifecycle_status = 'superseded', updated_at = ? WHERE id = ? """, (scan_started_at, reused_existing_occurrence_row["id"]), ) superseded_document_id = int(existing_document_row["id"]) refresh_source_backed_dataset_memberships_for_document(connection, superseded_document_id) refresh_document_from_occurrences(connection, superseded_document_id) existing_document_row = None reused_existing_occurrence_row = None extracted = apply_manual_locks(existing_document_row, extracted_payload) attachments = list(prepared_item.get("attachments") or []) if existing_document_row is None: if current_ingestion_batch is None: current_ingestion_batch = allocate_ingestion_batch_number(connection) control_number_batch = current_ingestion_batch control_number_family_sequence = reserve_control_number_family_sequence(connection, control_number_batch) control_number = format_control_number(control_number_batch, control_number_family_sequence) control_number_attachment_sequence = None else: control_number_batch = int(existing_document_row["control_number_batch"]) control_number_family_sequence = int(existing_document_row["control_number_family_sequence"]) control_number = str(existing_document_row["control_number"]) control_number_attachment_sequence = existing_document_row["control_number_attachment_sequence"] cleanup_document_artifacts(paths, connection, existing_document_row) document_id = upsert_document_row( connection, str(prepared_item["rel_path"]), None, existing_document_row, extracted, existing_occurrence_row=reused_existing_occurrence_row, file_name=str(prepared_item["file_name"]), parent_document_id=None, control_number=control_number, dataset_id=dataset_id, control_number_batch=control_number_batch, control_number_family_sequence=control_number_family_sequence, control_number_attachment_sequence=control_number_attachment_sequence, source_kind=source_kind, source_rel_path=source_rel_path, source_item_id=str(prepared_item["source_item_id"]), source_folder_path=source_folder_path, file_type_override=file_type_override, file_size_override=None, file_hash_override=file_hash, ingested_at_override=scan_started_at, last_seen_at_override=scan_started_at, updated_at_override=scan_started_at, ) replace_document_email_threading_row( connection, document_id=document_id, email_threading=extracted.get("email_threading"), ) replace_document_chat_threading_row( connection, document_id=document_id, chat_threading=extracted.get("chat_threading"), ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=existing_document_row, created_at=scan_started_at, ) preview_rows = write_preview_artifacts( paths, str(prepared_item["rel_path"]), list(extracted.get("preview_artifacts", [])), ) replace_document_related_rows( connection, document_id, extracted | {"file_name": str(prepared_item["file_name"])}, list(prepared_item.get("prepared_chunks") or []), preview_rows, ) ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=dataset_source_id, ) reconcile_attachment_documents( connection, paths, document_id, str(prepared_item["rel_path"]), control_number_batch, control_number_family_sequence, attachments, [(dataset_id, dataset_source_id)], ) if superseded_document_id is not None and superseded_document_id != document_id: refresh_source_backed_dataset_memberships_for_document(connection, superseded_document_id) refresh_document_from_occurrences(connection, superseded_document_id) result = { "action": "new" if existing_document_row is None else "updated", "current_ingestion_batch": current_ingestion_batch, "document_id": document_id, } if before_transaction_commit is not None: before_transaction_commit(connection, result) return result def commit_prepared_container_message( connection: sqlite3.Connection, paths: dict[str, Path], prepared_item: dict[str, object], existing_row: sqlite3.Row | None, existing_occurrence_row: sqlite3.Row | None, *, current_ingestion_batch: int | None, dataset_id: int, dataset_source_id: int | None, source_kind: str, source_rel_path: str, file_type_override: str, scan_started_at: str, before_transaction_commit=None, ) -> dict[str, object]: connection.execute("BEGIN") try: result = commit_prepared_container_message_in_transaction( connection, paths, prepared_item, existing_row, existing_occurrence_row, current_ingestion_batch=current_ingestion_batch, dataset_id=dataset_id, dataset_source_id=dataset_source_id, source_kind=source_kind, source_rel_path=source_rel_path, file_type_override=file_type_override, scan_started_at=scan_started_at, before_transaction_commit=before_transaction_commit, ) connection.commit() return result except Exception: connection.rollback() raise def ingest_container_source( connection: sqlite3.Connection, paths: dict[str, Path], path: Path, source_rel_path: str, *, source_kind: str, scan_hash_salt: str, dataset_name: str, iter_messages, normalize_message, file_type_override: str, source_scan_hash_override: str | None = None, staging_root: Path | None = None, ) -> dict[str, object]: source_scan_hash = normalize_whitespace(str(source_scan_hash_override or "")) or sha256_text( f"{scan_hash_salt}:{sha256_file(path) or ''}" ) transaction_was_open = connection.in_transaction dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=source_kind, source_locator=source_rel_path, dataset_name=dataset_name, ) if not transaction_was_open and connection.in_transaction: # Source-backed dataset repair may create or reattach dataset metadata # for legacy workspaces. Flush that implicit transaction before the # per-source BEGIN blocks below. connection.commit() existing_source = get_container_source_row(connection, source_kind, source_rel_path) file_size = file_size_bytes(path) file_mtime = file_mtime_timestamp(path) scan_started_at = next_monotonic_utc_timestamp( [ existing_source["last_scan_started_at"] if existing_source is not None else None, existing_source["last_scan_completed_at"] if existing_source is not None else None, ] ) if ( existing_source is not None and container_source_scan_completed(existing_source) and not container_documents_missing_text_revisions( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ): same_size = existing_source["file_size"] == file_size same_mtime = existing_source["file_mtime"] == file_mtime file_hash = source_scan_hash if ( same_size and same_mtime and existing_source["file_hash"] == source_scan_hash and not container_email_documents_missing_threading( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ): message_count = int(existing_source["message_count"] or 0) if message_count == 0: message_count = len( container_root_occurrence_rows_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ) connection.execute("BEGIN") try: mark_container_source_documents_active( connection, source_kind=source_kind, source_rel_path=source_rel_path, seen_at=scan_started_at, ) assign_dataset_to_container_documents( connection, source_kind=source_kind, source_rel_path=source_rel_path, dataset_id=dataset_id, dataset_source_id=dataset_source_id, ) write_container_source_scan_completed( connection, dataset_id=dataset_id, source_kind=source_kind, source_rel_path=source_rel_path, file_size=file_size, file_mtime=file_mtime, file_hash=file_hash, message_count=message_count, scan_started_at=scan_started_at, scan_completed_at=scan_started_at, ) connection.commit() except Exception: connection.rollback() raise return { "action": "skipped", "container_sources_skipped": 1, "container_messages_created": 0, "container_messages_updated": 0, "container_messages_deleted": 0, "container_prepare_ms": 0.0, "container_chunk_ms": 0.0, "container_prepare_wait_ms": 0.0, "container_commit_ms": 0.0, } if ( same_size and existing_source["file_hash"] and existing_source["file_hash"] == source_scan_hash and not container_email_documents_missing_threading( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ): message_count = int(existing_source["message_count"] or 0) if message_count == 0: message_count = len( container_root_occurrence_rows_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ) connection.execute("BEGIN") try: mark_container_source_documents_active( connection, source_kind=source_kind, source_rel_path=source_rel_path, seen_at=scan_started_at, ) assign_dataset_to_container_documents( connection, source_kind=source_kind, source_rel_path=source_rel_path, dataset_id=dataset_id, dataset_source_id=dataset_source_id, ) write_container_source_scan_completed( connection, dataset_id=dataset_id, source_kind=source_kind, source_rel_path=source_rel_path, file_size=file_size, file_mtime=file_mtime, file_hash=file_hash, message_count=message_count, scan_started_at=scan_started_at, scan_completed_at=scan_started_at, ) connection.commit() except Exception: connection.rollback() raise return { "action": "skipped", "container_sources_skipped": 1, "container_messages_created": 0, "container_messages_updated": 0, "container_messages_deleted": 0, "container_prepare_ms": 0.0, "container_chunk_ms": 0.0, "container_prepare_wait_ms": 0.0, "container_commit_ms": 0.0, } else: file_hash = source_scan_hash write_container_source_scan_started( connection, dataset_id=dataset_id, source_kind=source_kind, source_rel_path=source_rel_path, file_size=file_size, file_mtime=file_mtime, file_hash=file_hash, scan_started_at=scan_started_at, ) connection.commit() existing_entries_by_source_item = existing_container_entries_by_source_item( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) current_ingestion_batch: int | None = None container_messages_created = 0 container_messages_updated = 0 message_count = 0 container_prepare_ms = 0.0 container_chunk_ms = 0.0 container_prepare_wait_ms = 0.0 container_commit_ms = 0.0 for prepared_item, wait_ms in iter_prepared_container_message_items( source_kind=source_kind, source_rel_path=source_rel_path, raw_messages=iter_messages(path), normalize_message=normalize_message, staging_root=staging_root, ): container_prepare_wait_ms += wait_ms container_prepare_ms += float(prepared_item["prepare_ms"]) container_chunk_ms += float(prepared_item.get("prepare_chunk_ms") or 0.0) if prepared_item.get("skip"): continue message_count += 1 existing_entry = existing_entries_by_source_item.get(str(prepared_item["source_item_id"])) existing_row = ( existing_entry["document_row"] if isinstance(existing_entry, dict) else None ) existing_occurrence_row = ( existing_entry["occurrence_row"] if isinstance(existing_entry, dict) else None ) commit_started = time.perf_counter() commit_result = commit_prepared_container_message( connection, paths, prepared_item, existing_row, existing_occurrence_row, current_ingestion_batch=current_ingestion_batch, dataset_id=dataset_id, dataset_source_id=dataset_source_id, source_kind=source_kind, source_rel_path=source_rel_path, file_type_override=file_type_override, scan_started_at=scan_started_at, ) container_commit_ms += (time.perf_counter() - commit_started) * 1000.0 current_ingestion_batch = commit_result["current_ingestion_batch"] if str(commit_result["action"]) == "new": container_messages_created += 1 else: container_messages_updated += 1 connection.execute("BEGIN") try: container_messages_deleted = retire_unseen_container_messages( connection, paths, source_kind=source_kind, source_rel_path=source_rel_path, scan_started_at=scan_started_at, ) scan_completed_at = next_monotonic_utc_timestamp([scan_started_at]) write_container_source_scan_completed( connection, dataset_id=dataset_id, source_kind=source_kind, source_rel_path=source_rel_path, file_size=file_size, file_mtime=file_mtime, file_hash=file_hash, message_count=message_count, scan_started_at=scan_started_at, scan_completed_at=scan_completed_at, ) connection.commit() except Exception: connection.rollback() raise return { "action": "new" if existing_source is None else "updated", "container_sources_skipped": 0, "container_messages_created": container_messages_created, "container_messages_updated": container_messages_updated, "container_messages_deleted": container_messages_deleted, "container_prepare_ms": container_prepare_ms, "container_chunk_ms": container_chunk_ms, "container_prepare_wait_ms": container_prepare_wait_ms, "container_commit_ms": container_commit_ms, } def mark_missing_pst_documents( connection: sqlite3.Connection, scanned_source_rel_paths: set[str], scan_scope: dict[str, object] | None = None, ) -> tuple[int, int]: return mark_missing_container_documents( connection, source_kind=PST_SOURCE_KIND, scanned_source_rel_paths=scanned_source_rel_paths, scan_scope=scan_scope, ) def mark_missing_mbox_documents( connection: sqlite3.Connection, scanned_source_rel_paths: set[str], scan_scope: dict[str, object] | None = None, ) -> tuple[int, int]: return mark_missing_container_documents( connection, source_kind=MBOX_SOURCE_KIND, scanned_source_rel_paths=scanned_source_rel_paths, scan_scope=scan_scope, ) def ingest_pst_source( connection: sqlite3.Connection, paths: dict[str, Path], path: Path, source_rel_path: str, *, message_metadata_by_source_item: dict[str, dict[str, object]] | None = None, message_match_records: list[dict[str, object]] | None = None, message_sidecar_hash: str | None = None, staging_root: Path | None = None, ) -> dict[str, object]: normalized_message_metadata = { str(key): dict(value) for key, value in dict(message_metadata_by_source_item or {}).items() } normalized_message_match_records = [ dict(record) for record in list(message_match_records or []) if isinstance(record, dict) ] def normalize_enriched_pst_message( source_rel_path_for_message: str, message_dict: dict[str, object], ) -> dict[str, object] | None: normalized = normalize_pst_message(source_rel_path_for_message, message_dict) if normalized is None: return None message_metadata = select_pst_export_message_metadata( normalized, exact_metadata_by_source_item=normalized_message_metadata, message_match_records=normalized_message_match_records, ) if not message_metadata: return normalized enriched = dict(normalized) enriched["extracted"] = apply_pst_export_message_metadata( dict(normalized["extracted"]), message_metadata=message_metadata, identifier_scope=source_rel_path_for_message, ) enriched["file_hash"] = pst_export_enriched_message_file_hash( normalized.get("file_hash"), message_metadata=message_metadata, ) return enriched pst_scan_hash_override = ( sha256_json_value( { "pst_hash": sha256_file(path), "message_sidecar_hash": message_sidecar_hash, "sidecar_match_version": "pst-export-sidecar-v2", "source_rel_path": source_rel_path, } ) if normalize_whitespace(str(message_sidecar_hash or "")) else None ) # Salt the scan fingerprint so unchanged PSTs get one corrective reparse when # container-routing rules change (for example, when Teams/system folders are reclassified # or Teams participant identity extraction improves). result = ingest_container_source( connection, paths, path, source_rel_path, source_kind=PST_SOURCE_KIND, scan_hash_salt="pst-ingest-v5", dataset_name=pst_dataset_name(source_rel_path), iter_messages=iter_pst_messages, normalize_message=normalize_enriched_pst_message, file_type_override=PST_SOURCE_KIND, source_scan_hash_override=pst_scan_hash_override, staging_root=staging_root, ) return { "action": result["action"], "pst_sources_skipped": result["container_sources_skipped"], "pst_messages_created": result["container_messages_created"], "pst_messages_updated": result["container_messages_updated"], "pst_messages_deleted": result["container_messages_deleted"], "pst_prepare_ms": result["container_prepare_ms"], "pst_chunk_ms": result["container_chunk_ms"], "pst_prepare_wait_ms": result["container_prepare_wait_ms"], "pst_commit_ms": result["container_commit_ms"], } def ingest_mbox_source( connection: sqlite3.Connection, paths: dict[str, Path], path: Path, source_rel_path: str, staging_root: Path | None = None, ) -> dict[str, object]: result = ingest_container_source( connection, paths, path, source_rel_path, source_kind=MBOX_SOURCE_KIND, scan_hash_salt="mbox-ingest-v1", dataset_name=mbox_dataset_name(source_rel_path), iter_messages=iter_mbox_messages, normalize_message=normalize_mbox_message, file_type_override=MBOX_SOURCE_KIND, staging_root=staging_root, ) return { "action": result["action"], "mbox_sources_skipped": result["container_sources_skipped"], "mbox_messages_created": result["container_messages_created"], "mbox_messages_updated": result["container_messages_updated"], "mbox_messages_deleted": result["container_messages_deleted"], "mbox_prepare_ms": result["container_prepare_ms"], "mbox_chunk_ms": result["container_chunk_ms"], "mbox_prepare_wait_ms": result["container_prepare_wait_ms"], "mbox_commit_ms": result["container_commit_ms"], } def remove_auto_filesystem_dataset_membership( connection: sqlite3.Connection, *, document_id: int, ) -> None: filesystem_source_row = get_dataset_source_row( connection, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), ) if filesystem_source_row is None: return connection.execute( """ DELETE FROM dataset_documents WHERE document_id = ? AND dataset_source_id = ? """, (document_id, int(filesystem_source_row["id"])), ) def retire_standalone_filesystem_documents_by_rel_paths( connection: sqlite3.Connection, paths: dict[str, Path], *, rel_paths: set[str], ) -> int: normalized_rel_paths = sorted( { normalize_whitespace(str(rel_path or "")) for rel_path in rel_paths if normalize_whitespace(str(rel_path or "")) } ) if not normalized_rel_paths: return 0 placeholders = ", ".join("?" for _ in normalized_rel_paths) rows = connection.execute( f""" SELECT * FROM documents WHERE parent_document_id IS NULL AND rel_path IN ({placeholders}) AND COALESCE(source_kind, ?) = ? AND lifecycle_status != 'deleted' ORDER BY id ASC """, [*normalized_rel_paths, FILESYSTEM_SOURCE_KIND, FILESYSTEM_SOURCE_KIND], ).fetchall() retired = 0 now = utc_now() for row in rows: child_rows = connection.execute( """ SELECT * FROM documents WHERE parent_document_id = ? ORDER BY id ASC """, (row["id"],), ).fetchall() for child_row in child_rows: cleanup_document_artifacts(paths, connection, child_row) delete_document_related_rows(connection, int(child_row["id"])) cleanup_document_artifacts(paths, connection, row) delete_document_related_rows(connection, int(row["id"])) related_ids = [int(row["id"]), *[int(child_row["id"]) for child_row in child_rows]] related_placeholders = ", ".join("?" for _ in related_ids) connection.execute( f""" UPDATE documents SET lifecycle_status = 'deleted', updated_at = ? WHERE id IN ({related_placeholders}) """, [now, *related_ids], ) retired += 1 return retired def ingest_gmail_export_root( connection: sqlite3.Connection, paths: dict[str, Path], root: Path, descriptor: dict[str, object], allowed_file_types: set[str] | None = None, staging_root: Path | None = None, ) -> dict[str, object]: all_mbox_paths = [Path(path) for path in list(descriptor.get("mbox_paths") or [])] if not all_mbox_paths: return { "new": 0, "updated": 0, "failed": 0, "scanned_files": 0, "mbox_sources_skipped": 0, "mbox_messages_created": 0, "mbox_messages_updated": 0, "mbox_messages_deleted": 0, "gmail_linked_documents_created": 0, "gmail_linked_documents_updated": 0, "scanned_filesystem_rel_paths": [], "scanned_mbox_source_rel_paths": [], "failures": [], } include_mbox_sources = allowed_file_types is None or MBOX_SOURCE_KIND in allowed_file_types selected_drive_file_types = ( None if allowed_file_types is None else {file_type for file_type in allowed_file_types if file_type != MBOX_SOURCE_KIND} ) mbox_paths = list(all_mbox_paths) if include_mbox_sources else [] drive_documents: list[dict[str, object]] = [] for raw_drive_record in list(descriptor.get("drive_documents") or []): drive_record = dict(raw_drive_record) file_path_value = drive_record.get("file_path") if not isinstance(file_path_value, Path): continue if include_mbox_sources and list(drive_record.get("linked_message_ids") or []): continue if ( selected_drive_file_types is not None and normalize_extension(file_path_value) not in selected_drive_file_types ): continue drive_documents.append(drive_record) if not mbox_paths and not drive_documents: return { "new": 0, "updated": 0, "failed": 0, "scanned_files": 0, "mbox_sources_skipped": 0, "mbox_messages_created": 0, "mbox_messages_updated": 0, "mbox_messages_deleted": 0, "gmail_linked_documents_created": 0, "gmail_linked_documents_updated": 0, "scanned_filesystem_rel_paths": [], "scanned_mbox_source_rel_paths": [], "failures": [], } primary_source_rel_path = relative_document_path(root, all_mbox_paths[0]) transaction_was_open = connection.in_transaction dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=MBOX_SOURCE_KIND, source_locator=primary_source_rel_path, dataset_name=mbox_dataset_name(primary_source_rel_path), ) if not transaction_was_open and connection.in_transaction: # Gmail exports can repair legacy dataset-source rows before any of the # explicit per-source BEGIN blocks run. connection.commit() email_metadata_by_message_id = { str(key): dict(value) for key, value in dict(descriptor.get("email_metadata_by_message_id") or {}).items() } linked_drive_attachment_records_by_message_id = { str(key): [dict(item) for item in list(value)] for key, value in dict(descriptor.get("linked_drive_attachment_records_by_message_id") or {}).items() } linked_drive_records_by_message_id = { str(key): [dict(item) for item in list(value)] for key, value in dict(descriptor.get("linked_drive_records_by_message_id") or {}).items() } message_sidecar_hash = normalize_whitespace(str(descriptor.get("message_sidecar_hash") or "")) or None stats = { "new": 0, "updated": 0, "failed": 0, "scanned_files": 0, "mbox_sources_skipped": 0, "mbox_messages_created": 0, "mbox_messages_updated": 0, "mbox_messages_deleted": 0, "gmail_linked_documents_created": 0, "gmail_linked_documents_updated": 0, } failures: list[dict[str, str]] = [] scanned_filesystem_rel_paths: set[str] = set() scanned_mbox_source_rel_paths: set[str] = set() current_ingestion_batch: int | None = None gmail_staging_root = None if staging_root is not None: gmail_staging_root = ( Path(staging_root) / "gmail" / sanitize_storage_filename(relative_document_path(root, Path(descriptor["root"]))) ) for mbox_path in mbox_paths: source_rel_path = relative_document_path(root, mbox_path) scanned_mbox_source_rel_paths.add(source_rel_path) def normalize_gmail_message(source_rel_path_for_message: str, message_dict: dict[str, object]) -> dict[str, object]: normalized = normalize_mbox_message(source_rel_path_for_message, message_dict) message_id = gmail_normalized_message_lookup_key( message_dict.get("source_item_id") or normalized.get("source_item_id") ) linked_drive_attachment_records = ( list(linked_drive_attachment_records_by_message_id.get(message_id, [])) if message_id is not None else [] ) linked_drive_records = ( list(linked_drive_records_by_message_id.get(message_id, [])) if message_id is not None else [] ) extracted = apply_gmail_email_export_metadata( dict(normalized["extracted"]), message_metadata=(email_metadata_by_message_id.get(message_id) if message_id is not None else None), linked_drive_records=linked_drive_records, ) attachment_payloads = [ attachment for attachment in ( gmail_drive_attachment_payload(dict(record)) for record in linked_drive_attachment_records ) if attachment is not None ] if attachment_payloads: extracted["attachments"] = [*list(extracted.get("attachments") or []), *attachment_payloads] normalized["extracted"] = extracted normalized["file_hash"] = gmail_enriched_message_file_hash( normalized.get("file_hash"), message_metadata=(email_metadata_by_message_id.get(message_id) if message_id is not None else None), linked_drive_records=linked_drive_records, linked_drive_attachment_records=linked_drive_attachment_records, ) return normalized mbox_scan_hash_override = sha256_json_value( { "mbox_hash": sha256_file(mbox_path), "message_sidecar_hash": message_sidecar_hash, "source_rel_path": source_rel_path, } ) result = ingest_container_source( connection, paths, mbox_path, source_rel_path, source_kind=MBOX_SOURCE_KIND, scan_hash_salt="mbox-ingest-v2-gmail", dataset_name=mbox_dataset_name(source_rel_path), iter_messages=iter_mbox_messages, normalize_message=normalize_gmail_message, file_type_override=MBOX_SOURCE_KIND, source_scan_hash_override=mbox_scan_hash_override, staging_root=gmail_staging_root, ) if result["action"] == "new": stats["new"] += 1 elif result["action"] == "updated": stats["updated"] += 1 stats["scanned_files"] += 1 stats["mbox_sources_skipped"] += int(result["container_sources_skipped"]) stats["mbox_messages_created"] += int(result["container_messages_created"]) stats["mbox_messages_updated"] += int(result["container_messages_updated"]) stats["mbox_messages_deleted"] += int(result["container_messages_deleted"]) if include_mbox_sources: linked_drive_rel_paths = { relative_document_path(root, file_path) for records in linked_drive_attachment_records_by_message_id.values() for record in records if isinstance(file_path := record.get("file_path"), Path) and file_path.exists() } if linked_drive_rel_paths: connection.execute("BEGIN") try: retire_standalone_filesystem_documents_by_rel_paths( connection, paths, rel_paths=linked_drive_rel_paths, ) connection.commit() except Exception: connection.rollback() raise for drive_record in drive_documents: file_path_value = drive_record.get("file_path") if not isinstance(file_path_value, Path) or not file_path_value.exists(): continue rel_path = relative_document_path(root, file_path_value) scanned_filesystem_rel_paths.add(rel_path) existing_row = connection.execute( """ SELECT * FROM documents WHERE parent_document_id IS NULL AND rel_path = ? ORDER BY id ASC LIMIT 1 """, (rel_path,), ).fetchone() connection.execute("BEGIN") try: extracted_payload = extract_document(file_path_value, include_attachments=True) attachments = list(extracted_payload.get("attachments", [])) extracted_payload.pop("attachments", None) extracted_payload = apply_gmail_drive_export_metadata( dict(extracted_payload), drive_record=drive_record, ) extracted = apply_manual_locks(existing_row, extracted_payload) if existing_row is None: if current_ingestion_batch is None: current_ingestion_batch = allocate_ingestion_batch_number(connection) control_number_batch = current_ingestion_batch control_number_family_sequence = reserve_control_number_family_sequence(connection, control_number_batch) control_number = format_control_number(control_number_batch, control_number_family_sequence) control_number_attachment_sequence = None else: control_number_batch = int(existing_row["control_number_batch"]) control_number_family_sequence = int(existing_row["control_number_family_sequence"]) control_number = str(existing_row["control_number"]) control_number_attachment_sequence = existing_row["control_number_attachment_sequence"] cleanup_document_artifacts(paths, connection, existing_row) document_id = upsert_document_row( connection, rel_path, file_path_value, existing_row, extracted, file_name=file_path_value.name, parent_document_id=None, control_number=control_number, dataset_id=dataset_id, control_number_batch=control_number_batch, control_number_family_sequence=control_number_family_sequence, control_number_attachment_sequence=control_number_attachment_sequence, source_kind=FILESYSTEM_SOURCE_KIND, file_hash_override=gmail_drive_document_file_hash(file_path_value, drive_record), ) replace_document_email_threading_row( connection, document_id=document_id, email_threading=extracted.get("email_threading"), ) replace_document_chat_threading_row( connection, document_id=document_id, chat_threading=extracted.get("chat_threading"), ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=existing_row, ) preview_rows = write_preview_artifacts(paths, rel_path, list(extracted.get("preview_artifacts", []))) chunks = extracted_search_chunks(extracted) replace_document_related_rows( connection, document_id, extracted | {"file_name": file_path_value.name}, chunks, preview_rows, ) remove_auto_filesystem_dataset_membership(connection, document_id=document_id) ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=dataset_source_id, ) reconcile_attachment_documents( connection, paths, document_id, rel_path, control_number_batch, control_number_family_sequence, attachments, [(dataset_id, dataset_source_id)], ) connection.commit() stats["scanned_files"] += 1 if existing_row is None: stats["new"] += 1 stats["gmail_linked_documents_created"] += 1 else: stats["updated"] += 1 stats["gmail_linked_documents_updated"] += 1 except Exception as exc: connection.rollback() stats["failed"] += 1 failures.append( { "rel_path": rel_path, "error": f"{type(exc).__name__}: {exc}", } ) return { **stats, "scanned_filesystem_rel_paths": sorted(scanned_filesystem_rel_paths), "scanned_mbox_source_rel_paths": sorted(scanned_mbox_source_rel_paths), "failures": failures, } def upsert_production_row( connection: sqlite3.Connection, *, dataset_id: int | None, rel_root: str, production_name: str, metadata_load_rel_path: str, image_load_rel_path: str | None, source_type: str, ) -> int: now = utc_now() existing_row = connection.execute( """ SELECT id FROM productions WHERE rel_root = ? """, (rel_root,), ).fetchone() if existing_row is None: connection.execute( """ INSERT INTO productions ( dataset_id, rel_root, production_name, metadata_load_rel_path, image_load_rel_path, source_type, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, (dataset_id, rel_root, production_name, metadata_load_rel_path, image_load_rel_path, source_type, now, now), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) connection.execute( """ UPDATE productions SET dataset_id = ?, production_name = ?, metadata_load_rel_path = ?, image_load_rel_path = ?, source_type = ?, updated_at = ? WHERE id = ? """, (dataset_id, production_name, metadata_load_rel_path, image_load_rel_path, source_type, now, existing_row["id"]), ) return int(existing_row["id"]) def production_previewable_native(path: Path | None) -> Path | None: if path is None or not path.exists(): return None file_type = normalize_extension(path) if file_type == "pdf": return path if file_type in SUPPORTED_FILE_TYPES: return path return None def production_source_parts( workspace_root: Path, *, text_path: Path | None, image_paths: list[Path], native_path: Path | None, ) -> list[dict[str, object]]: parts: list[dict[str, object]] = [] if text_path is not None and text_path.exists(): parts.append( { "part_kind": "text", "rel_source_path": relative_document_path(workspace_root, text_path), "ordinal": 0, "label": "Linked text", "created_at": utc_now(), } ) for index, image_path in enumerate(image_paths, start=1): if not image_path.exists(): continue parts.append( { "part_kind": "image", "rel_source_path": relative_document_path(workspace_root, image_path), "ordinal": index, "label": f"Page {index}", "created_at": utc_now(), } ) if native_path is not None and native_path.exists(): parts.append( { "part_kind": "native", "rel_source_path": relative_document_path(workspace_root, native_path), "ordinal": 0, "label": native_path.name, "created_at": utc_now(), } ) return parts def build_production_extracted_payload( workspace_root: Path, *, production_name: str, control_number: str, begin_bates: str, end_bates: str, begin_attachment: str | None, end_attachment: str | None, text_path: Path | None, image_paths: list[Path], native_path: Path | None, preview_image_limit: int | None = None, preview_image_max_dimension: int | None = None, preview_image_refs: list[dict[str, object]] | None = None, embed_preview_images: bool = False, ) -> dict[str, object]: text_content = "" text_status = "empty" if text_path is not None and text_path.exists(): text_content, text_status, _ = decode_bytes(text_path.read_bytes()) text_content = normalize_whitespace(text_content) if not text_content: text_status = "empty" email_headers = extract_email_like_headers(text_content) author = email_headers.get("author") if email_headers else None recipients = email_headers.get("recipients") if email_headers else None subject = email_headers.get("subject") if email_headers else None participants = extract_email_chain_participants( text_content, [author, recipients] if email_headers else None, ) or extract_chat_participants(text_content) file_type = normalize_extension(native_path) if native_path is not None else "" preferred_native = production_previewable_native(native_path) image_content_type = infer_content_type_from_extension(normalize_extension(image_paths[0])) if image_paths else None fallback_content_type = infer_content_type_from_extension(file_type) or image_content_type or "E-Doc" content_type_path = text_path or preferred_native or native_path or (image_paths[0] if image_paths else Path(f"{control_number}.txt")) content_type = ( determine_content_type( content_type_path, text_content, email_headers=email_headers or None, explicit_content_type=fallback_content_type, ) or fallback_content_type ) page_images: list[dict[str, object]] = [] page_image_note = None if preview_image_refs is not None and not embed_preview_images: for ref in preview_image_refs: page_images.append( { "label": str(ref.get("label") or f"Page {ref.get('ordinal') or len(page_images) + 1}"), "src": str(ref.get("html_src") or ""), } ) if page_images: page_image_note = "Produced page previews are generated in resumable batches." else: preview_image_paths = image_paths normalized_preview_image_limit = None if preview_image_limit is None else max(0, int(preview_image_limit)) if normalized_preview_image_limit is not None: preview_image_paths = image_paths[:normalized_preview_image_limit] for index, image_path in enumerate(preview_image_paths, start=1): if not image_path.exists(): continue data_url = image_path_data_url(image_path, max_dimension=preview_image_max_dimension) if data_url is None: continue page_images.append({"label": f"Page {index}", "src": data_url}) if normalized_preview_image_limit is not None and len(image_paths) > normalized_preview_image_limit: page_image_note = ( f"Preview shows the first {len(preview_image_paths)} of {len(image_paths)} produced pages. " "All produced page files are still linked as source parts." ) resolved_title = (email_headers.get("title") if email_headers else None) or infer_production_title(control_number, text_content, native_path) preview_artifacts: list[dict[str, object]] = [] if preferred_native is None: preview_artifacts.append( { "file_name": f"{sanitize_storage_filename(control_number)}.html", "preview_type": "html", "label": "production", "ordinal": 0, "content": build_production_preview_html( document_title=resolved_title, control_number=control_number, production_name=production_name, begin_bates=begin_bates, end_bates=end_bates, begin_attachment=begin_attachment, end_attachment=end_attachment, text_content=text_content, page_images=page_images, page_image_note=page_image_note, ), } ) return { "page_count": len(image_paths) or None, "author": author, "content_type": content_type, "date_created": email_headers.get("date_created") if email_headers else None, "date_modified": None, "participants": participants, "title": resolved_title, "subject": subject, "recipients": recipients, "text_content": text_content, "text_status": text_status, "preview_artifacts": preview_artifacts, "preferred_native": preferred_native, } def production_document_file_size(text_path: Path | None, image_paths: list[Path], native_path: Path | None) -> int | None: total = 0 found = False for path in [text_path, native_path, *image_paths]: if path is None or not path.exists(): continue total += path.stat().st_size found = True return total if found else None def plan_production_record_work( workspace_root: Path, resolved_production_root: Path, signature: dict[str, object], metadata_rows: list[dict[str, str]], resolved_image_rows: list[dict[str, object]], ) -> tuple[list[dict[str, object]], set[str]]: plans: list[dict[str, object]] = [] seen_control_numbers: set[str] = set() for record in metadata_rows: begin_bates = str(record.get("begin_bates") or "").strip() end_bates = str(record.get("end_bates") or begin_bates).strip() if not begin_bates: continue control_number = begin_bates seen_control_numbers.add(control_number) text_path = resolve_production_source_path(workspace_root, resolved_production_root, record.get("text_path")) native_path = resolve_production_source_path(workspace_root, resolved_production_root, record.get("native_path")) matching_image_paths = [ Path(image_row["resolved_path"]) for image_row in resolved_image_rows if image_row.get("resolved_path") is not None and bates_inclusive_contains(begin_bates, end_bates, image_row["page_bates"]) ] plans.append( { "production_name": str(signature["production_name"]), "production_rel_root": str(signature["rel_root"]), "control_number": control_number, "begin_bates": begin_bates, "end_bates": end_bates, "begin_attachment": record.get("begin_attachment"), "end_attachment": record.get("end_attachment"), "text_path": text_path, "native_path": native_path, "matching_image_paths": matching_image_paths, "missing_linked_text": bool(record.get("text_path") and (text_path is None or not text_path.exists())), "missing_linked_images": bool(resolved_image_rows and not matching_image_paths), "missing_linked_natives": bool(record.get("native_path") and (native_path is None or not native_path.exists())), } ) return plans, seen_control_numbers def prepare_production_row_plan( workspace_root: Path, prepared_plan: dict[str, object], *, preview_image_limit: int | None = None, preview_image_max_dimension: int | None = None, preview_image_refs: list[dict[str, object]] | None = None, embed_preview_images: bool = False, ) -> dict[str, object]: prepared_item = dict(prepared_plan) prepare_started = time.perf_counter() try: text_path = Path(prepared_plan["text_path"]) if prepared_plan.get("text_path") is not None else None native_path = Path(prepared_plan["native_path"]) if prepared_plan.get("native_path") is not None else None matching_image_paths = [Path(path) for path in list(prepared_plan.get("matching_image_paths") or [])] available_text_path = text_path if text_path is not None and text_path.exists() else None available_native_path = native_path if native_path is not None and native_path.exists() else None extracted_payload = build_production_extracted_payload( workspace_root, production_name=str(prepared_plan["production_name"]), control_number=str(prepared_plan["control_number"]), begin_bates=str(prepared_plan["begin_bates"]), end_bates=str(prepared_plan["end_bates"]), begin_attachment=prepared_plan.get("begin_attachment"), end_attachment=prepared_plan.get("end_attachment"), text_path=available_text_path, image_paths=matching_image_paths, native_path=available_native_path, preview_image_limit=preview_image_limit, preview_image_max_dimension=preview_image_max_dimension, preview_image_refs=preview_image_refs, embed_preview_images=embed_preview_images, ) preferred_native = extracted_payload.pop("preferred_native", None) source_parts = production_source_parts( workspace_root, text_path=available_text_path, image_paths=matching_image_paths, native_path=available_native_path, ) preview_page_assets = ( production_preview_page_assets( preview_image_refs, max_dimension=preview_image_max_dimension, ) if preview_image_refs is not None and not embed_preview_images and preferred_native is None else [] ) chunk_started = time.perf_counter() prepared_chunks = chunk_text(str(extracted_payload.get("text_content") or "")) rel_path = production_logical_rel_path(str(prepared_plan["production_rel_root"]), str(prepared_plan["control_number"])).as_posix() file_name = ( (preferred_native.name if isinstance(preferred_native, Path) else None) or (available_native_path.name if available_native_path is not None else None) or f"{prepared_plan['control_number']}.production" ) preferred_source_path = ( preferred_native if isinstance(preferred_native, Path) else available_text_path ) prepared_item["extracted_payload"] = extracted_payload prepared_item["preferred_native"] = preferred_native prepared_item["preferred_source_path"] = preferred_source_path prepared_item["source_parts"] = source_parts prepared_item["preview_page_assets"] = preview_page_assets prepared_item["prepared_chunks"] = prepared_chunks prepared_item["prepare_chunk_ms"] = (time.perf_counter() - chunk_started) * 1000.0 prepared_item["rel_path"] = rel_path prepared_item["file_name"] = file_name prepared_item["available_text_path"] = available_text_path prepared_item["available_native_path"] = available_native_path prepared_item["matching_image_paths"] = matching_image_paths prepared_item["file_type_override"] = ( normalize_extension(preferred_native) if isinstance(preferred_native, Path) else (normalize_extension(native_path) if native_path is not None else None) ) prepared_item["file_size_override"] = production_document_file_size( available_text_path, matching_image_paths, available_native_path, ) prepared_item["file_hash_override"] = ( sha256_file(preferred_native) if isinstance(preferred_native, Path) else (sha256_file(available_text_path) if available_text_path is not None else None) ) prepared_item["prepare_error"] = None except Exception as exc: prepared_item["extracted_payload"] = None prepared_item["preferred_native"] = None prepared_item["preferred_source_path"] = None prepared_item["source_parts"] = [] prepared_item["preview_page_assets"] = [] prepared_item["prepared_chunks"] = [] prepared_item["prepare_chunk_ms"] = 0.0 prepared_item["rel_path"] = production_logical_rel_path( str(prepared_plan["production_rel_root"]), str(prepared_plan["control_number"]), ).as_posix() prepared_item["file_name"] = f"{prepared_plan['control_number']}.production" prepared_item["available_text_path"] = None prepared_item["available_native_path"] = None prepared_item["matching_image_paths"] = [] prepared_item["file_type_override"] = None prepared_item["file_size_override"] = None prepared_item["file_hash_override"] = None prepared_item["prepare_error"] = f"{type(exc).__name__}: {exc}" prepared_item["prepare_ms"] = (time.perf_counter() - prepare_started) * 1000.0 return prepared_item def production_expected_preview_artifacts_present( paths: dict[str, Path], connection: sqlite3.Connection, *, document_id: int, expected_html_preview_paths: list[str], expected_page_preview_paths: list[str], ) -> bool: expected_paths = [ normalize_whitespace(str(path)) for path in [*expected_html_preview_paths, *expected_page_preview_paths] if normalize_whitespace(str(path)) ] if not expected_paths: return True rows = connection.execute( """ SELECT rel_preview_path FROM document_previews WHERE document_id = ? """, (document_id,), ).fetchall() known_paths = {str(row["rel_preview_path"] or "") for row in rows} for rel_preview_path in expected_paths: if rel_preview_path not in known_paths: return False if not (paths["state_dir"] / rel_preview_path).exists(): return False return True def production_plan_matches_existing_document( paths: dict[str, Path], connection: sqlite3.Connection, workspace_root: Path, existing_row: sqlite3.Row | None, prepared_plan: dict[str, object], *, production_id: int, preview_image_max_dimension: int | None = None, embed_preview_images: bool = False, ) -> bool: if existing_row is None or existing_row["lifecycle_status"] != "active": return False try: text_path = Path(prepared_plan["text_path"]) if prepared_plan.get("text_path") is not None else None native_path = Path(prepared_plan["native_path"]) if prepared_plan.get("native_path") is not None else None matching_image_paths = [Path(path) for path in list(prepared_plan.get("matching_image_paths") or [])] available_text_path = text_path if text_path is not None and text_path.exists() else None available_native_path = native_path if native_path is not None and native_path.exists() else None rel_path = production_logical_rel_path( str(prepared_plan["production_rel_root"]), str(prepared_plan["control_number"]), ).as_posix() preview_image_refs = production_preview_page_asset_refs( rel_path, str(prepared_plan["control_number"]), matching_image_paths, ) extracted_payload = build_production_extracted_payload( workspace_root, production_name=str(prepared_plan["production_name"]), control_number=str(prepared_plan["control_number"]), begin_bates=str(prepared_plan["begin_bates"]), end_bates=str(prepared_plan["end_bates"]), begin_attachment=prepared_plan.get("begin_attachment"), end_attachment=prepared_plan.get("end_attachment"), text_path=available_text_path, image_paths=matching_image_paths, native_path=available_native_path, preview_image_max_dimension=preview_image_max_dimension, preview_image_refs=preview_image_refs, embed_preview_images=embed_preview_images, ) preferred_native = extracted_payload.get("preferred_native") source_parts = production_source_parts( workspace_root, text_path=available_text_path, image_paths=matching_image_paths, native_path=available_native_path, ) file_name = ( (preferred_native.name if isinstance(preferred_native, Path) else None) or (available_native_path.name if available_native_path is not None else None) or f"{prepared_plan['control_number']}.production" ) desired_signature = production_row_signature( existing_row, rel_path=rel_path, file_name=file_name, source_kind=PRODUCTION_SOURCE_KIND, production_id=production_id, begin_bates=str(prepared_plan["begin_bates"]), end_bates=str(prepared_plan["end_bates"]), begin_attachment=prepared_plan.get("begin_attachment"), end_attachment=prepared_plan.get("end_attachment"), extracted=extracted_payload, source_parts=source_parts, ) if existing_production_row_signature(connection, existing_row) != desired_signature: return False preview_base = preview_base_path_for_rel_path(rel_path) expected_html_preview_paths = [ (preview_base / str(artifact["file_name"])).as_posix() for artifact in list(extracted_payload.get("preview_artifacts") or []) if isinstance(artifact, dict) and artifact.get("file_name") ] expected_page_preview_paths = ( [] if embed_preview_images else [ str(ref["rel_preview_path"]) for ref in preview_image_refs if ref.get("rel_preview_path") ] ) if embed_preview_images: stale_page_preview_row = connection.execute( """ SELECT 1 FROM document_previews WHERE document_id = ? AND preview_type = 'image' LIMIT 1 """, (int(existing_row["id"]),), ).fetchone() if stale_page_preview_row is not None: return False return production_expected_preview_artifacts_present( paths, connection, document_id=int(existing_row["id"]), expected_html_preview_paths=expected_html_preview_paths, expected_page_preview_paths=expected_page_preview_paths, ) except Exception: return False def iter_prepared_production_row_plans( workspace_root: Path, production_row_plans: list[dict[str, object]], staging_root: Path | None = None, ) -> Iterator[tuple[dict[str, object], float]]: effective_staging_root = staging_root if staging_root is not None and production_row_plans: effective_staging_root = ( Path(staging_root) / "production" / sanitize_storage_filename(str(production_row_plans[0]["production_rel_root"])) ) yield from iter_staged_prepared_items( production_row_plans, prepare_item=lambda plan: prepare_production_row_plan( workspace_root, plan, preview_image_max_dimension=( INGEST_V2_PRODUCTION_PREVIEW_IMAGE_MAX_DIMENSION if not INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES else None ), preview_image_refs=( production_preview_page_asset_refs( production_logical_rel_path( str(plan["production_rel_root"]), str(plan["control_number"]), ).as_posix(), str(plan["control_number"]), [ Path(str(path)) for path in list(plan.get("matching_image_paths") or []) if path ], ) if not INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES else None ), embed_preview_images=INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES, ), config_benchmark_name="ingest_production_prepare_config", queue_done_benchmark_name="ingest_production_prepare_queue_done", spill_subdir_name="prepared-production", staging_root=effective_staging_root, ) def commit_prepared_production_row( connection: sqlite3.Connection, paths: dict[str, Path], existing_row: sqlite3.Row | None, prepared_item: dict[str, object], *, dataset_id: int, dataset_source_id: int, production_id: int, before_transaction_commit=None, ) -> dict[str, object]: control_number = str(prepared_item["control_number"]) prepare_error = prepared_item.get("prepare_error") if prepare_error: return { "action": "failed", "control_number": control_number, "error": str(prepare_error), "page_images_linked": 0, } connection.execute("BEGIN") try: existing_signature = existing_production_row_signature(connection, existing_row) extracted = apply_manual_locks(existing_row, dict(prepared_item["extracted_payload"] or {})) desired_signature = production_row_signature( existing_row, rel_path=str(prepared_item["rel_path"]), file_name=str(prepared_item["file_name"]), source_kind=PRODUCTION_SOURCE_KIND, production_id=production_id, begin_bates=str(prepared_item["begin_bates"]), end_bates=str(prepared_item["end_bates"]), begin_attachment=prepared_item.get("begin_attachment"), end_attachment=prepared_item.get("end_attachment"), extracted=extracted, source_parts=list(prepared_item.get("source_parts", [])), ) if existing_row is not None: cleanup_document_artifacts(paths, connection, existing_row) document_id = upsert_document_row( connection, str(prepared_item["rel_path"]), ( prepared_item["preferred_source_path"] if isinstance(prepared_item.get("preferred_source_path"), Path) else None ), existing_row, extracted, file_name=str(prepared_item["file_name"]), parent_document_id=None, control_number=control_number, dataset_id=dataset_id, control_number_batch=None, control_number_family_sequence=None, control_number_attachment_sequence=None, source_kind=PRODUCTION_SOURCE_KIND, production_id=production_id, begin_bates=str(prepared_item["begin_bates"]), end_bates=str(prepared_item["end_bates"]), begin_attachment=prepared_item.get("begin_attachment"), end_attachment=prepared_item.get("end_attachment"), file_type_override=prepared_item.get("file_type_override"), file_size_override=prepared_item.get("file_size_override"), file_hash_override=prepared_item.get("file_hash_override"), ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=existing_row, ) ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=dataset_source_id, ) preview_rows = write_preview_artifacts(paths, str(prepared_item["rel_path"]), list(extracted.get("preview_artifacts", []))) preview_rows.extend(write_preview_page_assets(paths, list(prepared_item.get("preview_page_assets", [])))) replace_document_related_rows( connection, document_id, extracted | {"file_name": str(prepared_item["file_name"])}, list(prepared_item.get("prepared_chunks", [])), preview_rows, ) replace_document_source_parts(connection, document_id, list(prepared_item.get("source_parts", []))) if existing_row is None: action = "created" elif existing_row["lifecycle_status"] == "active" and existing_signature == desired_signature: action = "unchanged" else: action = "updated" result = { "action": action, "control_number": control_number, "page_images_linked": len(list(prepared_item.get("matching_image_paths", []))), "document_id": document_id, "production_id": production_id, } if before_transaction_commit is not None: before_transaction_commit(connection, result) connection.commit() return result except Exception as exc: connection.rollback() return { "action": "failed", "control_number": control_number, "error": f"{type(exc).__name__}: {exc}", "page_images_linked": 0, } def update_production_family_relationships(connection: sqlite3.Connection, production_id: int) -> int: rows = connection.execute( """ SELECT id, control_number, begin_bates, end_bates, begin_attachment, end_attachment, parent_document_id FROM documents WHERE production_id = ? AND lifecycle_status != 'deleted' ORDER BY begin_bates ASC, control_number ASC, id ASC """, (production_id,), ).fetchall() parents: list[sqlite3.Row] = [] for row in rows: if row["begin_attachment"] and row["end_attachment"] and row["control_number"] == row["begin_attachment"]: parents.append(row) updated = 0 for row in rows: desired_parent_id = None row_begin = row["begin_bates"] or row["control_number"] row_end = row["end_bates"] or row["control_number"] for parent_row in parents: if int(parent_row["id"]) == int(row["id"]): continue if bates_ranges_overlap(parent_row["begin_attachment"], parent_row["end_attachment"], row_begin, row_end): desired_parent_id = int(parent_row["id"]) break if row["parent_document_id"] == desired_parent_id: continue connection.execute( "UPDATE documents SET parent_document_id = ?, updated_at = ? WHERE id = ?", (desired_parent_id, utc_now(), row["id"]), ) updated += 1 return updated def ingest_resolved_production_root( connection: sqlite3.Connection, paths: dict[str, Path], workspace_root: Path, resolved_production_root: Path, staging_root: Path | None = None, ) -> dict[str, object]: workspace_root = workspace_root.resolve() resolved_production_root = resolved_production_root.resolve() signature = production_signature_for_root(workspace_root, resolved_production_root) if signature is None: raise RetrieverError(f"Path does not look like a supported processed production: {resolved_production_root}") metadata_load_path = Path(signature["metadata_load_path"]) image_load_path = Path(signature["image_load_path"]) if signature["image_load_path"] is not None else None metadata = parse_production_metadata_load(metadata_load_path) image_rows = parse_production_image_load(image_load_path) dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=PRODUCTION_SOURCE_KIND, source_locator=str(signature["rel_root"]), dataset_name=production_dataset_name(str(signature["rel_root"]), str(signature["production_name"])), ) production_id = upsert_production_row( connection, dataset_id=dataset_id, rel_root=str(signature["rel_root"]), production_name=str(signature["production_name"]), metadata_load_rel_path=relative_document_path(workspace_root, metadata_load_path), image_load_rel_path=relative_document_path(workspace_root, image_load_path) if image_load_path is not None else None, source_type=str(signature["source_type"]), ) connection.commit() existing_rows = connection.execute( """ SELECT * FROM documents WHERE production_id = ? """, (production_id,), ).fetchall() existing_by_control_number = {str(row["control_number"]): row for row in existing_rows if row["control_number"]} resolved_image_rows: list[dict[str, object]] = [] for image_row in image_rows: resolved_path = resolve_production_source_path(workspace_root, resolved_production_root, image_row["image_path"]) resolved_image_rows.append({**image_row, "resolved_path": resolved_path}) production_row_plans, seen_control_numbers = plan_production_record_work( workspace_root, resolved_production_root, signature, list(metadata["rows"]), resolved_image_rows, ) stats: dict[str, int] = { "created": 0, "updated": 0, "unchanged": 0, "retired": 0, "families_reconstructed": 0, "page_images_linked": 0, "docs_missing_linked_text": 0, "docs_missing_linked_images": 0, "docs_missing_linked_natives": 0, } failures: list[dict[str, str]] = [] stats["docs_missing_linked_text"] = sum(int(plan["missing_linked_text"]) for plan in production_row_plans) stats["docs_missing_linked_images"] = sum(int(plan["missing_linked_images"]) for plan in production_row_plans) stats["docs_missing_linked_natives"] = sum(int(plan["missing_linked_natives"]) for plan in production_row_plans) prepare_ms = 0.0 prepare_chunk_ms = 0.0 prepare_wait_ms = 0.0 commit_ms = 0.0 row_loop_started = time.perf_counter() for prepared_item, wait_ms in iter_prepared_production_row_plans( workspace_root, production_row_plans, staging_root=staging_root, ): prepare_wait_ms += wait_ms prepare_ms += float(prepared_item.get("prepare_ms") or 0.0) prepare_chunk_ms += float(prepared_item.get("prepare_chunk_ms") or 0.0) control_number = str(prepared_item["control_number"]) existing_row = existing_by_control_number.get(control_number) commit_started = time.perf_counter() commit_result = commit_prepared_production_row( connection, paths, existing_row, prepared_item, dataset_id=dataset_id, dataset_source_id=dataset_source_id, production_id=production_id, ) commit_ms += (time.perf_counter() - commit_started) * 1000.0 action = str(commit_result["action"]) if action == "failed": failures.append( { "control_number": control_number, "error": str(commit_result.get("error") or "Unknown production ingest failure."), } ) continue stats[action] += 1 stats["page_images_linked"] += int(commit_result["page_images_linked"]) benchmark_mark( "ingest_production_rows_done", row_count=len(production_row_plans), row_loop_ms=round((time.perf_counter() - row_loop_started) * 1000.0, 3), prepare_ms=round(prepare_ms, 3), prepare_chunk_ms=round(prepare_chunk_ms, 3), prepare_wait_ms=round(prepare_wait_ms, 3), commit_ms=round(commit_ms, 3), created=stats["created"], updated=stats["updated"], unchanged=stats["unchanged"], failed=len(failures), ) retire_started = time.perf_counter() for row in existing_rows: control_number = str(row["control_number"] or "") if control_number and control_number in seen_control_numbers: continue connection.execute("BEGIN") try: cleanup_document_artifacts(paths, connection, row) delete_document_related_rows(connection, int(row["id"])) connection.execute( """ UPDATE documents SET lifecycle_status = 'deleted', parent_document_id = NULL, updated_at = ? WHERE id = ? """, (utc_now(), row["id"]), ) connection.commit() stats["retired"] += 1 except Exception: connection.rollback() raise benchmark_mark( "ingest_production_retire_done", retire_ms=round((time.perf_counter() - retire_started) * 1000.0, 3), retired=stats["retired"], ) finalize_started = time.perf_counter() connection.execute("BEGIN") try: parent_link_updates = update_production_family_relationships(connection, production_id) attachment_preview_updates = 0 preview_document_rows = connection.execute( """ SELECT DISTINCT documents.id FROM documents JOIN document_previews ON document_previews.document_id = documents.id WHERE documents.production_id = ? AND documents.lifecycle_status != 'deleted' AND document_previews.preview_type = 'html' ORDER BY documents.id ASC """, (production_id,), ).fetchall() for preview_document_row in preview_document_rows: attachment_preview_updates += sync_document_attachment_preview_links( connection, paths, int(preview_document_row["id"]), ) connection.commit() except Exception: connection.rollback() raise stats["families_reconstructed"] = len( connection.execute( """ SELECT id FROM documents WHERE production_id = ? AND parent_document_id IS NOT NULL AND lifecycle_status != 'deleted' """, (production_id,), ).fetchall() ) stats["parent_link_updates"] = parent_link_updates stats["attachment_preview_updates"] = attachment_preview_updates benchmark_mark( "ingest_production_finalize_done", finalize_ms=round((time.perf_counter() - finalize_started) * 1000.0, 3), parent_link_updates=parent_link_updates, attachment_preview_updates=attachment_preview_updates, families_reconstructed=stats["families_reconstructed"], ) return { "status": "ok", "workspace_root": str(workspace_root), "production_root": str(resolved_production_root), "production_rel_root": str(signature["rel_root"]), "production_name": str(signature["production_name"]), "production_id": production_id, "metadata_load_rel_path": relative_document_path(workspace_root, metadata_load_path), "image_load_rel_path": relative_document_path(workspace_root, image_load_path) if image_load_path is not None else None, "tool_version": TOOL_VERSION, "schema_version": SCHEMA_VERSION, "failures": failures, **stats, } def inspect_custom_fields_registry(connection: sqlite3.Connection) -> dict[str, object]: columns = table_info(connection, "documents") column_map = {row["name"]: row for row in columns} actual_custom_fields = [ name for name in column_map if name not in BUILTIN_FIELD_TYPES and name not in INTERNAL_DOCUMENT_COLUMNS and name not in {LEGACY_METADATA_LOCKS_COLUMN} ] registry_rows = connection.execute( """ SELECT field_name, field_type, instruction, created_at FROM custom_fields_registry ORDER BY field_name ASC """ ).fetchall() registry_fields = {row["field_name"]: row for row in registry_rows} missing_registry = sorted(name for name in actual_custom_fields if name not in registry_fields) orphaned_registry = sorted(name for name in registry_fields if name not in column_map) return { "actual_custom_fields": sorted(actual_custom_fields), "missing_registry": missing_registry, "orphaned_registry": orphaned_registry, } def reconcile_custom_fields_registry(connection: sqlite3.Connection, repair: bool) -> dict[str, object]: status = inspect_custom_fields_registry(connection) repaired: list[str] = [] if repair: now = utc_now() for field_name in status["missing_registry"]: sqlite_type = next( (row["type"] for row in table_info(connection, "documents") if row["name"] == field_name), "", ) connection.execute( """ INSERT INTO custom_fields_registry (field_name, field_type, instruction, created_at) VALUES (?, ?, ?, ?) ON CONFLICT(field_name) DO NOTHING """, (field_name, infer_registry_field_type(sqlite_type), None, now), ) repaired.append(field_name) if repaired: connection.commit() status = inspect_custom_fields_registry(connection) status["repaired_registry"] = repaired return status def merge_legacy_field_locks(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") if MANUAL_FIELD_LOCKS_COLUMN not in columns or LEGACY_METADATA_LOCKS_COLUMN not in columns: return 0 rows = connection.execute( f""" SELECT id, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} AS manual_locks, {quote_identifier(LEGACY_METADATA_LOCKS_COLUMN)} AS legacy_locks FROM documents """ ).fetchall() merged_count = 0 for row in rows: manual_locks = normalize_string_list(row["manual_locks"]) legacy_locks = normalize_string_list(row["legacy_locks"]) if not legacy_locks: continue merged = list(manual_locks) for field_name in legacy_locks: if field_name not in merged: merged.append(field_name) if merged == manual_locks: continue connection.execute( f""" UPDATE documents SET {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} = ? WHERE id = ? """, (json.dumps(merged), row["id"]), ) merged_count += 1 return merged_count def backfill_content_type(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") if "content_type" not in columns: return 0 rows = connection.execute( """ SELECT id, file_type, content_type FROM documents WHERE content_type IS NULL OR TRIM(content_type) = '' """ ).fetchall() updated = 0 for row in rows: inferred = infer_content_type_from_extension((row["file_type"] or "").lower()) if not inferred: continue connection.execute( "UPDATE documents SET content_type = ? WHERE id = ?", (inferred, row["id"]), ) updated += 1 return updated def backfill_child_document_kinds(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") required = {"parent_document_id", "child_document_kind"} if not required.issubset(columns): return 0 cursor = connection.execute( """ UPDATE documents SET child_document_kind = ? WHERE parent_document_id IS NOT NULL AND (child_document_kind IS NULL OR TRIM(child_document_kind) = '') """, (CHILD_DOCUMENT_KIND_ATTACHMENT,), ) return int(cursor.rowcount or 0) def backfill_conversation_assignment_modes(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") if "conversation_assignment_mode" not in columns: return 0 cursor = connection.execute( """ UPDATE documents SET conversation_assignment_mode = ? WHERE conversation_assignment_mode IS NULL OR TRIM(conversation_assignment_mode) = '' """, (CONVERSATION_ASSIGNMENT_MODE_AUTO,), ) return int(cursor.rowcount or 0) def backfill_custodian(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") required = {"custodian", "source_kind", "source_rel_path", "parent_document_id"} if not required.issubset(columns): return 0 rows = connection.execute( """ SELECT id, source_kind, source_rel_path, parent_document_id FROM documents WHERE custodian IS NULL OR TRIM(custodian) = '' ORDER BY CASE WHEN parent_document_id IS NULL THEN 0 ELSE 1 END ASC, id ASC """ ).fetchall() updated = 0 for row in rows: parent_custodian = None parent_document_id = row["parent_document_id"] if parent_document_id is not None: parent_row = connection.execute( "SELECT custodian FROM documents WHERE id = ?", (parent_document_id,), ).fetchone() if parent_row is not None: parent_custodian = parent_row["custodian"] custodian = infer_source_custodian( source_kind=row["source_kind"], source_rel_path=row["source_rel_path"], parent_custodian=parent_custodian, ) if not custodian: continue connection.execute( "UPDATE documents SET custodian = ? WHERE id = ?", (custodian, row["id"]), ) updated += 1 return updated def backfill_dataset_ids(connection: sqlite3.Connection, root: Path | None = None) -> int: if not table_exists(connection, "datasets") or not table_exists(connection, "documents"): return 0 productions_by_id: dict[int, sqlite3.Row] = {} if table_exists(connection, "productions"): production_rows = connection.execute( """ SELECT id, dataset_id, rel_root, production_name FROM productions ORDER BY id ASC """ ).fetchall() for row in production_rows: productions_by_id[int(row["id"])] = row updated = 0 if productions_by_id: for row in productions_by_id.values(): if row["dataset_id"] is not None: continue dataset_id = ensure_dataset_row( connection, source_kind=PRODUCTION_SOURCE_KIND, dataset_locator=str(row["rel_root"]), dataset_name=production_dataset_name(str(row["rel_root"]), str(row["production_name"] or "")), ) connection.execute( "UPDATE productions SET dataset_id = ? WHERE id = ?", (dataset_id, row["id"]), ) updated += 1 if table_exists(connection, "container_sources"): container_rows = connection.execute( """ SELECT id, dataset_id, source_kind, source_rel_path FROM container_sources ORDER BY id ASC """ ).fetchall() for row in container_rows: if row["dataset_id"] is not None: continue source_kind = str(row["source_kind"] or "") source_rel_path = str(row["source_rel_path"] or "") if not source_kind or not source_rel_path: continue dataset_name = ( pst_dataset_name(source_rel_path) if source_kind == PST_SOURCE_KIND else mbox_dataset_name(source_rel_path) if source_kind == MBOX_SOURCE_KIND else source_rel_path ) dataset_id = ensure_dataset_row( connection, source_kind=source_kind, dataset_locator=source_rel_path, dataset_name=dataset_name, ) connection.execute( "UPDATE container_sources SET dataset_id = ? WHERE id = ?", (dataset_id, row["id"]), ) updated += 1 document_columns = table_columns(connection, "documents") required_document_columns = {"id", "dataset_id", "parent_document_id", "source_kind", "source_rel_path", "production_id"} if not required_document_columns.issubset(document_columns): return updated filesystem_dataset_id: int | None = None rows = connection.execute( """ SELECT id, dataset_id, parent_document_id, source_kind, source_rel_path, production_id FROM documents WHERE dataset_id IS NULL ORDER BY CASE WHEN parent_document_id IS NULL THEN 0 ELSE 1 END ASC, id ASC """ ).fetchall() for row in rows: dataset_id: int | None = None parent_document_id = row["parent_document_id"] if parent_document_id is not None: parent_row = connection.execute( "SELECT dataset_id FROM documents WHERE id = ?", (parent_document_id,), ).fetchone() if parent_row is not None and parent_row["dataset_id"] is not None: dataset_id = int(parent_row["dataset_id"]) if dataset_id is None and row["production_id"] is not None: production_row = productions_by_id.get(int(row["production_id"])) if production_row is not None: if production_row["dataset_id"] is None: dataset_id = ensure_dataset_row( connection, source_kind=PRODUCTION_SOURCE_KIND, dataset_locator=str(production_row["rel_root"]), dataset_name=production_dataset_name( str(production_row["rel_root"]), str(production_row["production_name"] or ""), ), ) connection.execute( "UPDATE productions SET dataset_id = ? WHERE id = ?", (dataset_id, production_row["id"]), ) productions_by_id[int(production_row["id"])] = connection.execute( "SELECT id, dataset_id, rel_root, production_name FROM productions WHERE id = ?", (production_row["id"],), ).fetchone() updated += 1 else: dataset_id = int(production_row["dataset_id"]) source_kind = normalize_whitespace(str(row["source_kind"] or "")).lower() source_rel_path = normalize_whitespace(str(row["source_rel_path"] or "")) if dataset_id is None and source_kind in {PST_SOURCE_KIND, MBOX_SOURCE_KIND} and source_rel_path: dataset_id = ensure_dataset_row( connection, source_kind=source_kind, dataset_locator=source_rel_path, dataset_name=( pst_dataset_name(source_rel_path) if source_kind == PST_SOURCE_KIND else mbox_dataset_name(source_rel_path) ), ) if dataset_id is None: if filesystem_dataset_id is None: filesystem_dataset_id = ensure_dataset_row( connection, source_kind=FILESYSTEM_SOURCE_KIND, dataset_locator=filesystem_dataset_locator(), dataset_name=filesystem_dataset_name(root), ) dataset_id = filesystem_dataset_id connection.execute( "UPDATE documents SET dataset_id = ? WHERE id = ?", (dataset_id, row["id"]), ) updated += 1 return updated def backfill_dataset_memberships(connection: sqlite3.Connection) -> int: if not table_exists(connection, "dataset_documents") or not table_exists(connection, "datasets"): return 0 updated = 0 dataset_rows = connection.execute( """ SELECT id, source_kind, dataset_locator FROM datasets ORDER BY id ASC """ ).fetchall() for row in dataset_rows: source_kind = normalize_whitespace(str(row["source_kind"] or "")).lower() source_locator = normalize_whitespace(str(row["dataset_locator"] or "")) if not source_kind or not source_locator or source_kind == MANUAL_DATASET_SOURCE_KIND: continue ensure_dataset_source_row( connection, dataset_id=int(row["id"]), source_kind=source_kind, source_locator=source_locator, ) existing_membership_count = int( connection.execute("SELECT COUNT(*) AS count FROM dataset_documents").fetchone()["count"] or 0 ) if existing_membership_count > 0: return updated productions_by_id: dict[int, sqlite3.Row] = {} if table_exists(connection, "productions"): for row in connection.execute( """ SELECT id, dataset_id, rel_root FROM productions ORDER BY id ASC """ ).fetchall(): productions_by_id[int(row["id"])] = row document_columns = table_columns(connection, "documents") required_document_columns = {"id", "dataset_id", "parent_document_id", "source_kind", "source_rel_path", "production_id"} if not required_document_columns.issubset(document_columns): return updated rows = connection.execute( """ SELECT id, dataset_id, parent_document_id, source_kind, source_rel_path, production_id FROM documents WHERE dataset_id IS NOT NULL ORDER BY CASE WHEN parent_document_id IS NULL THEN 0 ELSE 1 END ASC, id ASC """ ).fetchall() for row in rows: dataset_id = int(row["dataset_id"]) document_id = int(row["id"]) source_membership_ids: list[int] = [] if row["parent_document_id"] is not None: parent_memberships = connection.execute( """ SELECT dataset_source_id FROM dataset_documents WHERE dataset_id = ? AND document_id = ? AND dataset_source_id IS NOT NULL ORDER BY dataset_source_id ASC """, (dataset_id, int(row["parent_document_id"])), ).fetchall() source_membership_ids = [int(item["dataset_source_id"]) for item in parent_memberships] if not source_membership_ids and row["production_id"] is not None: production_row = productions_by_id.get(int(row["production_id"])) if production_row is not None: dataset_source_row = get_dataset_source_row( connection, source_kind=PRODUCTION_SOURCE_KIND, source_locator=str(production_row["rel_root"]), ) if dataset_source_row is not None and int(dataset_source_row["dataset_id"]) == dataset_id: source_membership_ids.append(int(dataset_source_row["id"])) source_kind = normalize_whitespace(str(row["source_kind"] or "")).lower() source_rel_path = normalize_whitespace(str(row["source_rel_path"] or "")) if not source_membership_ids and source_kind in {PST_SOURCE_KIND, MBOX_SOURCE_KIND} and source_rel_path: dataset_source_row = get_dataset_source_row( connection, source_kind=source_kind, source_locator=source_rel_path, ) if dataset_source_row is not None and int(dataset_source_row["dataset_id"]) == dataset_id: source_membership_ids.append(int(dataset_source_row["id"])) if not source_membership_ids and source_kind in {FILESYSTEM_SOURCE_KIND, EMAIL_ATTACHMENT_SOURCE_KIND}: dataset_source_row = get_dataset_source_row( connection, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), ) if dataset_source_row is not None and int(dataset_source_row["dataset_id"]) == dataset_id: source_membership_ids.append(int(dataset_source_row["id"])) if source_membership_ids: for dataset_source_id in source_membership_ids: ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=dataset_source_id, ) updated += 1 continue ensure_dataset_document_membership( connection, dataset_id=dataset_id, document_id=document_id, dataset_source_id=None, ) updated += 1 return updated def backfill_occurrence_dataset_source_ids(connection: sqlite3.Connection) -> int: if not table_exists(connection, "document_occurrences") or not table_exists(connection, "dataset_sources"): return 0 columns = table_columns(connection, "document_occurrences") if "dataset_source_id" not in columns: return 0 rows = connection.execute( """ SELECT * FROM document_occurrences WHERE dataset_source_id IS NULL ORDER BY id ASC """ ).fetchall() updated = 0 for row in rows: dataset_source_row = dataset_source_row_for_occurrence(connection, row) if dataset_source_row is None: continue connection.execute( """ UPDATE document_occurrences SET dataset_source_id = ?, updated_at = COALESCE(updated_at, ?) WHERE id = ? """, (int(dataset_source_row["id"]), utc_now(), int(row["id"])), ) updated += 1 return updated def backfill_dataset_name_normalized(connection: sqlite3.Connection) -> int: if not table_exists(connection, "datasets"): return 0 columns = table_columns(connection, "datasets") if "dataset_name" not in columns or "dataset_name_normalized" not in columns: return 0 rows = connection.execute( """ SELECT id, dataset_name, dataset_name_normalized FROM datasets ORDER BY id ASC """ ).fetchall() updated = 0 for row in rows: dataset_name = normalized_dataset_name_or_default(str(row["dataset_name"] or "")) normalized_name = normalize_dataset_name_for_compare(dataset_name) if ( normalize_inline_whitespace(str(row["dataset_name"] or "")) == dataset_name and normalize_inline_whitespace(str(row["dataset_name_normalized"] or "")) == normalized_name ): continue connection.execute( """ UPDATE datasets SET dataset_name = ?, dataset_name_normalized = ?, updated_at = ? WHERE id = ? """, (dataset_name, normalized_name, utc_now(), int(row["id"])), ) updated += 1 return updated def merge_dataset_identity_duplicates(connection: sqlite3.Connection) -> int: if not table_exists(connection, "datasets"): return 0 rows = connection.execute( """ SELECT id, source_kind, dataset_locator FROM datasets ORDER BY LOWER(source_kind) ASC, dataset_locator ASC, id ASC """ ).fetchall() ids_by_identity: dict[tuple[str, str], list[int]] = defaultdict(list) for row in rows: identity = ( normalize_whitespace(str(row["source_kind"] or "")).lower(), normalize_whitespace(str(row["dataset_locator"] or "")), ) ids_by_identity[identity].append(int(row["id"])) merged = 0 for duplicate_ids in ids_by_identity.values(): if len(duplicate_ids) < 2: continue keep_id = min(duplicate_ids) for drop_id in sorted(dataset_id for dataset_id in duplicate_ids if dataset_id != keep_id): source_rows = connection.execute( """ SELECT id, source_kind, source_locator, created_at, updated_at FROM dataset_sources WHERE dataset_id = ? ORDER BY id ASC """, (drop_id,), ).fetchall() source_id_map: dict[int, int] = {} for source_row in source_rows: new_source_id = ensure_dataset_source_row( connection, dataset_id=keep_id, source_kind=str(source_row["source_kind"]), source_locator=str(source_row["source_locator"]), ) source_id_map[int(source_row["id"])] = new_source_id membership_rows = connection.execute( """ SELECT id, document_id, dataset_source_id, created_at, updated_at FROM dataset_documents WHERE dataset_id = ? ORDER BY id ASC """, (drop_id,), ).fetchall() for membership_row in membership_rows: dataset_source_id = membership_row["dataset_source_id"] remapped_source_id = ( source_id_map.get(int(dataset_source_id)) if dataset_source_id is not None else None ) connection.execute( """ INSERT OR IGNORE INTO dataset_documents ( dataset_id, document_id, dataset_source_id, created_at, updated_at ) VALUES (?, ?, ?, ?, ?) """, ( keep_id, int(membership_row["document_id"]), remapped_source_id, membership_row["created_at"] or utc_now(), membership_row["updated_at"] or utc_now(), ), ) connection.execute( "DELETE FROM dataset_documents WHERE id = ?", (int(membership_row["id"]),), ) connection.execute("UPDATE documents SET dataset_id = ? WHERE dataset_id = ?", (keep_id, drop_id)) connection.execute("UPDATE productions SET dataset_id = ? WHERE dataset_id = ?", (keep_id, drop_id)) connection.execute("UPDATE container_sources SET dataset_id = ? WHERE dataset_id = ?", (keep_id, drop_id)) connection.execute("DELETE FROM dataset_sources WHERE dataset_id = ?", (drop_id,)) connection.execute("DELETE FROM datasets WHERE id = ?", (drop_id,)) merged += 1 return merged def suffix_dataset_name_collisions(connection: sqlite3.Connection) -> int: if not table_exists(connection, "datasets"): return 0 rows = connection.execute( """ SELECT id, dataset_name FROM datasets ORDER BY id ASC """ ).fetchall() renamed = 0 used_names: set[str] = set() for row in rows: dataset_id = int(row["id"]) base_name = normalized_dataset_name_or_default(str(row["dataset_name"] or "")) candidate = base_name suffix = 2 normalized_candidate = normalize_dataset_name_for_compare(candidate) while normalized_candidate in used_names: candidate = f"{base_name}_{suffix}" normalized_candidate = normalize_dataset_name_for_compare(candidate) suffix += 1 used_names.add(normalized_candidate) if candidate == base_name and normalize_inline_whitespace(str(row["dataset_name"] or "")) == base_name: connection.execute( """ UPDATE datasets SET dataset_name_normalized = ?, updated_at = ? WHERE id = ? """, (normalized_candidate, utc_now(), dataset_id), ) continue connection.execute( """ UPDATE datasets SET dataset_name = ?, dataset_name_normalized = ?, updated_at = ? WHERE id = ? """, (candidate, normalized_candidate, utc_now(), dataset_id), ) if candidate != base_name or normalize_inline_whitespace(str(row["dataset_name"] or "")) != base_name: renamed += 1 return renamed def ensure_control_number_batch_row( connection: sqlite3.Connection, batch_number: int, next_family_sequence: int, ) -> int: row = connection.execute( """ SELECT next_family_sequence FROM control_number_batches WHERE batch_number = ? """, (batch_number,), ).fetchone() now = utc_now() if row is None: connection.execute( """ INSERT INTO control_number_batches (batch_number, next_family_sequence, created_at, updated_at) VALUES (?, ?, ?, ?) """, (batch_number, next_family_sequence, now, now), ) return next_family_sequence normalized_next = max(int(row["next_family_sequence"]), next_family_sequence) if normalized_next != int(row["next_family_sequence"]): connection.execute( """ UPDATE control_number_batches SET next_family_sequence = ?, updated_at = ? WHERE batch_number = ? """, (normalized_next, now, batch_number), ) return normalized_next def backfill_control_number_batches(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") required = {"control_number_batch", "control_number_family_sequence"} if not required.issubset(columns): return 0 rows = connection.execute( """ SELECT control_number_batch, MAX(control_number_family_sequence) AS max_family_sequence FROM documents WHERE control_number_batch IS NOT NULL AND control_number_family_sequence IS NOT NULL GROUP BY control_number_batch """ ).fetchall() updated = 0 for row in rows: before = connection.execute( """ SELECT next_family_sequence FROM control_number_batches WHERE batch_number = ? """, (row["control_number_batch"],), ).fetchone() next_sequence = int(row["max_family_sequence"] or 0) + 1 ensure_control_number_batch_row(connection, int(row["control_number_batch"]), next_sequence) if before is None or int(before["next_family_sequence"]) != next_sequence: updated += 1 return updated def backfill_document_occurrences(connection: sqlite3.Connection) -> int: if not table_exists(connection, "documents") or not table_exists(connection, "document_occurrences"): return 0 missing_document_rows = connection.execute( """ SELECT * FROM documents WHERE NOT EXISTS ( SELECT 1 FROM document_occurrences WHERE document_occurrences.document_id = documents.id ) AND canonical_status != ? ORDER BY CASE WHEN parent_document_id IS NULL THEN 0 ELSE 1 END ASC, id ASC """, (CANONICAL_STATUS_MERGED,), ).fetchall() preview_counts = { int(row["document_id"]): int(row["count"] or 0) for row in connection.execute( """ SELECT document_id, COUNT(*) AS count FROM document_previews WHERE document_id IN ( SELECT id FROM documents WHERE NOT EXISTS ( SELECT 1 FROM document_occurrences WHERE document_occurrences.document_id = documents.id ) ) GROUP BY document_id """ ).fetchall() } if missing_document_rows else {} inserted = 0 for row in missing_document_rows: document_id = int(row["id"]) legacy_custodian = None if "custodian" in row.keys(): legacy_custodian = row["custodian"] elif "custodians_json" in row.keys(): custodian_values = parse_document_custodians_json(row["custodians_json"]) legacy_custodian = custodian_values[0] if len(custodian_values) == 1 else None occurrence_id = upsert_document_occurrence( connection, document_id=document_id, existing_occurrence_id=None, parent_occurrence_id=None, occurrence_control_number=row["control_number"], source_kind=row["source_kind"], source_rel_path=row["source_rel_path"] or row["rel_path"], source_item_id=row["source_item_id"], source_folder_path=row["source_folder_path"], production_id=row["production_id"], begin_bates=row["begin_bates"], end_bates=row["end_bates"], begin_attachment=row["begin_attachment"], end_attachment=row["end_attachment"], rel_path=str(row["rel_path"]), file_name=str(row["file_name"]), file_type=row["file_type"], mime_type=None, file_size=row["file_size"], file_hash=row["file_hash"], custodian=legacy_custodian, fs_created_at=None, fs_modified_at=None, extracted={ "author": row["author"], "title": row["title"], "subject": row["subject"], "participants": row["participants"], "recipients": row["recipients"], "date_created": row["date_created"], "date_modified": row["date_modified"], "content_type": row["content_type"], }, has_preview=preview_counts.get(document_id, 0) > 0, text_status=str(row["text_status"] or "ok"), ingested_at=str(row["ingested_at"] or row["updated_at"] or utc_now()), last_seen_at=str(row["last_seen_at"] or row["ingested_at"] or row["updated_at"] or utc_now()), updated_at=str(row["updated_at"] or row["last_seen_at"] or row["ingested_at"] or utc_now()), ) inserted += 1 child_rows = connection.execute( """ SELECT child.id, child.parent_document_id, child_occurrence.id AS child_occurrence_id FROM documents child JOIN document_occurrences child_occurrence ON child_occurrence.document_id = child.id WHERE child.parent_document_id IS NOT NULL AND child_occurrence.parent_occurrence_id IS NULL ORDER BY child.id ASC, child_occurrence.id ASC """ ).fetchall() for row in child_rows: parent_occurrence = connection.execute( """ SELECT id FROM document_occurrences WHERE document_id = ? ORDER BY id ASC LIMIT 1 """, (int(row["parent_document_id"]),), ).fetchone() if parent_occurrence is None: continue connection.execute( """ UPDATE document_occurrences SET parent_occurrence_id = ? WHERE id = ? AND parent_occurrence_id IS NULL """, (int(parent_occurrence["id"]), int(row["child_occurrence_id"])), ) return inserted def drop_document_field_locks(connection: sqlite3.Connection, field_name: str) -> int: if not table_exists(connection, "documents"): return 0 rows = connection.execute( f""" SELECT id, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} AS locks_json FROM documents WHERE {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} LIKE ? """, (f'%"{field_name}"%',), ).fetchall() updated = 0 for row in rows: locks = normalize_string_list(row["locks_json"]) filtered_locks = [lock_name for lock_name in locks if lock_name != field_name] if filtered_locks == locks: continue connection.execute( f""" UPDATE documents SET {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} = ? WHERE id = ? """, (json.dumps(filtered_locks), int(row["id"])), ) updated += 1 return updated def remove_documents_custodian_column(connection: sqlite3.Connection) -> bool: if "custodian" not in table_columns(connection, "documents"): return False connection.execute("ALTER TABLE documents DROP COLUMN custodian") return True def backfill_document_dedupe_keys(connection: sqlite3.Connection) -> int: if not table_exists(connection, "documents") or not table_exists(connection, "document_dedupe_keys"): return 0 before = connection.total_changes rows = connection.execute( """ WITH desired_keys AS ( SELECT file_hash, MAX(id) AS document_id FROM documents WHERE parent_document_id IS NULL AND file_hash IS NOT NULL AND TRIM(file_hash) != '' AND lifecycle_status != 'deleted' AND LOWER(COALESCE(source_kind, '')) IN (?, ?, ?) GROUP BY file_hash ) SELECT desired_keys.document_id AS id, desired_keys.file_hash FROM desired_keys LEFT JOIN document_dedupe_keys existing_key ON existing_key.basis = 'file_hash' AND existing_key.key_value = desired_keys.file_hash WHERE existing_key.document_id IS NULL OR existing_key.document_id != desired_keys.document_id ORDER BY desired_keys.document_id ASC """, (FILESYSTEM_SOURCE_KIND, PST_SOURCE_KIND, MBOX_SOURCE_KIND), ).fetchall() for row in rows: bind_document_dedupe_key( connection, basis="file_hash", key_value=str(row["file_hash"]), document_id=int(row["id"]), ) return connection.total_changes - before def refresh_documents_from_occurrences(connection: sqlite3.Connection) -> int: if not table_exists(connection, "documents") or not table_exists(connection, "document_occurrences"): return 0 rows = connection.execute( """ SELECT id FROM documents ORDER BY CASE WHEN parent_document_id IS NULL THEN 0 ELSE 1 END ASC, id ASC """ ).fetchall() for row in rows: refresh_document_from_occurrences(connection, int(row["id"])) return len(rows) def should_refresh_document_occurrence_caches( *, prior_schema_version: int | None, backfilled_occurrence_dataset_source_ids: int, backfilled_document_occurrences: int, ) -> bool: # A full occurrence refresh also rebuilds document entity links; keep it off # the read-only command path once the workspace schema is current. if prior_schema_version != SCHEMA_VERSION: return True return backfilled_occurrence_dataset_source_ids > 0 or backfilled_document_occurrences > 0 def backfill_control_numbers(connection: sqlite3.Connection) -> int: columns = table_columns(connection, "documents") required = { "control_number", "parent_document_id", "control_number_batch", "control_number_family_sequence", "control_number_attachment_sequence", "source_kind", } if not required.issubset(columns): return 0 backfill_control_number_batches(connection) updated = 0 top_level_rows = connection.execute( """ SELECT id, control_number FROM documents WHERE parent_document_id IS NULL AND COALESCE(source_kind, ?) != ? AND ( control_number IS NULL OR TRIM(control_number) = '' OR control_number_batch IS NULL OR control_number_family_sequence IS NULL ) ORDER BY id ASC """ , (FILESYSTEM_SOURCE_KIND, PRODUCTION_SOURCE_KIND)).fetchall() if top_level_rows: next_family_sequence = ensure_control_number_batch_row(connection, 1, 1) for row in top_level_rows: parsed_identity = parse_control_number(row["control_number"]) if parsed_identity is not None and parsed_identity[2] is None: batch_number, family_sequence, _ = parsed_identity control_number = str(row["control_number"]).strip() ensure_control_number_batch_row(connection, batch_number, family_sequence + 1) if batch_number == 1: next_family_sequence = max(next_family_sequence, family_sequence + 1) else: batch_number = 1 family_sequence = next_family_sequence control_number = format_control_number(batch_number, family_sequence) next_family_sequence += 1 connection.execute( """ UPDATE documents SET control_number = ?, control_number_batch = ?, control_number_family_sequence = ?, control_number_attachment_sequence = NULL WHERE id = ? """, (control_number, batch_number, family_sequence, row["id"]), ) updated += 1 ensure_control_number_batch_row(connection, 1, next_family_sequence) parent_identities = { int(row["id"]): (int(row["control_number_batch"]), int(row["control_number_family_sequence"])) for row in connection.execute( """ SELECT id, control_number_batch, control_number_family_sequence FROM documents WHERE parent_document_id IS NULL AND COALESCE(source_kind, ?) != ? AND control_number_batch IS NOT NULL AND control_number_family_sequence IS NOT NULL """ , (FILESYSTEM_SOURCE_KIND, PRODUCTION_SOURCE_KIND)).fetchall() } attachment_counters: dict[int, int] = {} child_rows = connection.execute( """ SELECT id, parent_document_id, control_number FROM documents WHERE parent_document_id IS NOT NULL AND COALESCE(source_kind, ?) != ? AND ( control_number IS NULL OR TRIM(control_number) = '' OR control_number_batch IS NULL OR control_number_family_sequence IS NULL OR control_number_attachment_sequence IS NULL ) ORDER BY parent_document_id ASC, id ASC """ , (EMAIL_ATTACHMENT_SOURCE_KIND, PRODUCTION_SOURCE_KIND)).fetchall() for row in child_rows: parent_document_id = int(row["parent_document_id"]) parent_identity = parent_identities.get(parent_document_id) if parent_identity is None: continue parsed_identity = parse_control_number(row["control_number"]) batch_number, family_sequence = parent_identity if parsed_identity is not None and parsed_identity[2] is not None: _, _, attachment_sequence = parsed_identity else: if parent_document_id not in attachment_counters: max_sequence = connection.execute( """ SELECT MAX(control_number_attachment_sequence) AS max_attachment_sequence FROM documents WHERE parent_document_id = ? """, (parent_document_id,), ).fetchone() attachment_counters[parent_document_id] = int(max_sequence["max_attachment_sequence"] or 0) attachment_counters[parent_document_id] += 1 attachment_sequence = attachment_counters[parent_document_id] connection.execute( """ UPDATE documents SET control_number = ?, control_number_batch = ?, control_number_family_sequence = ?, control_number_attachment_sequence = ? WHERE id = ? """, ( format_control_number(batch_number, family_sequence, attachment_sequence), batch_number, family_sequence, attachment_sequence, row["id"], ), ) updated += 1 return updated def allocate_ingestion_batch_number(connection: sqlite3.Connection) -> int: row = connection.execute("SELECT MAX(batch_number) AS max_batch_number FROM control_number_batches").fetchone() batch_number = int(row["max_batch_number"] or 0) + 1 ensure_control_number_batch_row(connection, batch_number, 1) return batch_number def reserve_control_number_family_sequence(connection: sqlite3.Connection, batch_number: int) -> int: next_family_sequence = ensure_control_number_batch_row(connection, batch_number, 1) connection.execute( """ UPDATE control_number_batches SET next_family_sequence = ?, updated_at = ? WHERE batch_number = ? """, (next_family_sequence + 1, utc_now(), batch_number), ) return next_family_sequence def next_attachment_sequence(connection: sqlite3.Connection, parent_document_id: int) -> int: row = connection.execute( """ SELECT MAX(control_number_attachment_sequence) AS max_attachment_sequence FROM documents WHERE parent_document_id = ? """, (parent_document_id,), ).fetchone() return int(row["max_attachment_sequence"] or 0) + 1 def backfill_internal_rel_path_prefix(connection: sqlite3.Connection) -> int: """Rewrite synthetic rel_paths that still use the legacy ``.retriever/`` prefix. Container-derived messages, production logical documents, and attachment blobs used to be stored with a leading ``.retriever/`` so that ``/`` would happen to resolve to the real file under the state directory. That made it look like the workspace's opaque state directory contained indexed documents, which confused scans and searches. The canonical synthetic prefix is now ``_retriever/`` and path resolution translates that back to the state directory explicitly. """ if not table_exists(connection, "documents"): return 0 cursor = connection.execute( """ UPDATE documents SET rel_path = '_retriever/' || substr(rel_path, length('.retriever/') + 1) WHERE rel_path LIKE '.retriever/%' """ ) return int(cursor.rowcount or 0) def ensure_documents_fts(connection: sqlite3.Connection) -> bool: expected_columns = {"document_id", "file_name", "title", "subject", "author", "custodian", "participants", "recipients"} existing_columns = table_columns(connection, "documents_fts") if existing_columns == expected_columns: return False connection.execute("DROP TABLE IF EXISTS documents_fts") connection.execute( """ CREATE VIRTUAL TABLE documents_fts USING fts5( document_id UNINDEXED, file_name, title, subject, author, custodian, participants, recipients ) """ ) rows = connection.execute( """ SELECT id, file_name, title, subject, author, custodians_json, participants, recipients FROM documents """ ).fetchall() if rows: connection.executemany( """ INSERT INTO documents_fts (document_id, file_name, title, subject, author, custodian, participants, recipients) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, [ ( row["id"], row["file_name"], row["title"], row["subject"], row["author"], document_custodian_display_text_from_row(row), row["participants"], row["recipients"], ) for row in rows ], ) return True def apply_schema(connection: sqlite3.Connection, root: Path | None = None) -> dict[str, object]: prior_schema_version: int | None = None if table_exists(connection, "workspace_meta"): workspace_meta_row = connection.execute( """ SELECT schema_version FROM workspace_meta WHERE id = 1 """ ).fetchone() if workspace_meta_row is not None and workspace_meta_row["schema_version"] is not None: prior_schema_version = int(workspace_meta_row["schema_version"]) rename_table_if_needed(connection, "display_id_batches", "control_number_batches") for statement in SCHEMA_STATEMENTS: connection.execute(statement) migrated_export_work_items_nullable_document_id = ensure_export_work_items_nullable_document_id(connection) if table_exists(connection, "documents"): rename_column_if_needed(connection, "documents", "display_id", "control_number") rename_column_if_needed(connection, "documents", "display_batch", "control_number_batch") rename_column_if_needed(connection, "documents", "display_family_sequence", "control_number_family_sequence") rename_column_if_needed(connection, "documents", "display_attachment_sequence", "control_number_attachment_sequence") ensure_column(connection, "documents", f"{MANUAL_FIELD_LOCKS_COLUMN} TEXT NOT NULL DEFAULT '[]'") ensure_column(connection, "documents", "canonical_kind TEXT NOT NULL DEFAULT 'unknown'") ensure_column(connection, "documents", "canonical_status TEXT NOT NULL DEFAULT 'active'") ensure_column(connection, "documents", "merged_into_document_id INTEGER REFERENCES documents(id) ON DELETE SET NULL") ensure_column(connection, "documents", "content_type TEXT") ensure_column(connection, "documents", "custodians_json TEXT NOT NULL DEFAULT '[]'") ensure_column(connection, "documents", "participants TEXT") ensure_column(connection, "documents", "control_number TEXT") ensure_column(connection, "documents", "conversation_id INTEGER REFERENCES conversations(id) ON DELETE SET NULL") ensure_column(connection, "documents", f"conversation_assignment_mode TEXT NOT NULL DEFAULT '{CONVERSATION_ASSIGNMENT_MODE_AUTO}'") ensure_column(connection, "documents", "dataset_id INTEGER REFERENCES datasets(id) ON DELETE SET NULL") ensure_column(connection, "documents", "parent_document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE") ensure_column(connection, "documents", "child_document_kind TEXT") ensure_column(connection, "documents", "source_kind TEXT") ensure_column(connection, "documents", "source_rel_path TEXT") ensure_column(connection, "documents", "source_item_id TEXT") ensure_column(connection, "documents", "root_message_key TEXT") ensure_column(connection, "documents", "source_folder_path TEXT") ensure_column(connection, "documents", "production_id INTEGER REFERENCES productions(id) ON DELETE SET NULL") ensure_column(connection, "documents", "begin_bates TEXT") ensure_column(connection, "documents", "end_bates TEXT") ensure_column(connection, "documents", "begin_attachment TEXT") ensure_column(connection, "documents", "end_attachment TEXT") ensure_column(connection, "documents", "control_number_batch INTEGER") ensure_column(connection, "documents", "control_number_family_sequence INTEGER") ensure_column(connection, "documents", "control_number_attachment_sequence INTEGER") ensure_column(connection, "documents", "source_text_revision_id INTEGER REFERENCES text_revisions(id) ON DELETE SET NULL") ensure_column(connection, "documents", "active_search_text_revision_id INTEGER REFERENCES text_revisions(id) ON DELETE SET NULL") ensure_column(connection, "documents", "active_text_source_kind TEXT") ensure_column(connection, "documents", "active_text_language TEXT") ensure_column(connection, "documents", "active_text_quality_score REAL") ensure_column(connection, "document_occurrences", "mime_type TEXT") ensure_column(connection, "document_occurrences", "dataset_source_id INTEGER REFERENCES dataset_sources(id) ON DELETE SET NULL") ensure_column(connection, "document_occurrences", "fs_created_at TEXT") ensure_column(connection, "document_occurrences", "fs_modified_at TEXT") ensure_column(connection, "document_occurrences", "has_preview INTEGER NOT NULL DEFAULT 0") ensure_column(connection, "document_occurrences", "extracted_author TEXT") ensure_column(connection, "document_occurrences", "extracted_title TEXT") ensure_column(connection, "document_occurrences", "extracted_subject TEXT") ensure_column(connection, "document_occurrences", "extracted_participants TEXT") ensure_column(connection, "document_occurrences", "extracted_recipients TEXT") ensure_column(connection, "document_occurrences", "extracted_doc_authored_at TEXT") ensure_column(connection, "document_occurrences", "extracted_doc_modified_at TEXT") ensure_column(connection, "document_occurrences", "extracted_content_type TEXT") ensure_column(connection, "document_occurrences", "extracted_kind TEXT") ensure_column(connection, "document_occurrences", "entity_hints_json TEXT NOT NULL DEFAULT '{}'") ensure_column(connection, "document_email_threading", "message_id TEXT") ensure_column(connection, "document_email_threading", "in_reply_to TEXT") ensure_column(connection, "document_email_threading", "references_json TEXT NOT NULL DEFAULT '[]'") ensure_column(connection, "document_email_threading", "conversation_index TEXT") ensure_column(connection, "document_email_threading", "conversation_topic TEXT") ensure_column(connection, "document_email_threading", "normalized_subject TEXT") ensure_column(connection, "document_email_threading", "updated_at TEXT NOT NULL DEFAULT ''") ensure_column(connection, "document_chat_threading", "thread_id TEXT") ensure_column(connection, "document_chat_threading", "message_id TEXT") ensure_column(connection, "document_chat_threading", "parent_message_id TEXT") ensure_column(connection, "document_chat_threading", "thread_type TEXT") ensure_column(connection, "document_chat_threading", "participants_json TEXT NOT NULL DEFAULT '[]'") ensure_column(connection, "document_chat_threading", "updated_at TEXT NOT NULL DEFAULT ''") ensure_column(connection, "document_previews", "target_fragment TEXT") ensure_column(connection, "job_versions", "capability TEXT NOT NULL DEFAULT ''") ensure_column(connection, "runs", "activation_policy TEXT NOT NULL DEFAULT 'manual'") ensure_column(connection, "run_items", "result_id INTEGER REFERENCES results(id) ON DELETE SET NULL") ensure_column(connection, "run_items", "page_number INTEGER") ensure_column(connection, "run_items", "input_artifact_rel_path TEXT") ensure_column(connection, "run_items", "claimed_by TEXT") ensure_column(connection, "run_items", "claimed_at TEXT") ensure_column(connection, "run_items", "last_heartbeat_at TEXT") ensure_column(connection, "productions", "dataset_id INTEGER REFERENCES datasets(id) ON DELETE SET NULL") ensure_column(connection, "container_sources", "dataset_id INTEGER REFERENCES datasets(id) ON DELETE SET NULL") ensure_column(connection, "datasets", "dataset_name_normalized TEXT") ensure_column(connection, "datasets", "allow_auto_merge INTEGER NOT NULL DEFAULT 1") ensure_column(connection, "datasets", "email_auto_merge INTEGER NOT NULL DEFAULT 1") ensure_column(connection, "datasets", "handle_auto_merge INTEGER NOT NULL DEFAULT 1") ensure_column(connection, "datasets", "phone_auto_merge INTEGER NOT NULL DEFAULT 0") ensure_column(connection, "datasets", "name_auto_merge INTEGER NOT NULL DEFAULT 0") ensure_column(connection, "datasets", "external_id_auto_merge_names_json TEXT NOT NULL DEFAULT '[]'") backfilled_legacy_control_number = backfill_legacy_column( connection, "documents", "display_id", "control_number", treat_blank_as_missing=True, ) backfilled_legacy_control_number_batch = backfill_legacy_column( connection, "documents", "display_batch", "control_number_batch", ) backfilled_legacy_control_number_family_sequence = backfill_legacy_column( connection, "documents", "display_family_sequence", "control_number_family_sequence", ) backfilled_legacy_control_number_attachment_sequence = backfill_legacy_column( connection, "documents", "display_attachment_sequence", "control_number_attachment_sequence", ) if table_exists(connection, "job_versions") and table_exists(connection, "jobs"): blank_capability_rows = connection.execute( """ SELECT jv.id, j.job_kind FROM job_versions jv JOIN jobs j ON j.id = jv.job_id WHERE jv.capability IS NULL OR TRIM(jv.capability) = '' """ ).fetchall() for row in blank_capability_rows: job_kind = normalize_whitespace(str(row["job_kind"] or "")).lower() try: capability = default_job_capability_for_kind(job_kind) except RetrieverError: capability = "text_structured" connection.execute( "UPDATE job_versions SET capability = ? WHERE id = ?", (capability, int(row["id"])), ) connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_conversation_id ON documents(conversation_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_conversation_assignment_mode ON documents(conversation_assignment_mode)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_parent_document_id ON documents(parent_document_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_child_document_kind ON documents(child_document_kind)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_dataset_id ON documents(dataset_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_canonical_status ON documents(canonical_status)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_merged_into_document_id ON documents(merged_into_document_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_canonical_status_content_hash ON documents(canonical_status, content_hash)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_source_kind ON documents(source_kind)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_source_rel_path ON documents(source_rel_path)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_root_message_key ON documents(root_message_key)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_production_id ON documents(production_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_begin_bates ON documents(begin_bates)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_end_bates ON documents(end_bates)") connection.execute("CREATE INDEX IF NOT EXISTS idx_documents_source_text_revision_id ON documents(source_text_revision_id)") connection.execute( "CREATE INDEX IF NOT EXISTS idx_documents_active_search_text_revision_id ON documents(active_search_text_revision_id)" ) connection.execute("CREATE INDEX IF NOT EXISTS idx_document_occurrences_document_id ON document_occurrences(document_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_occurrences_document_status_ingested ON document_occurrences(document_id, lifecycle_status, ingested_at)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_occurrences_document_kind_status ON document_occurrences(document_id, extracted_kind, text_status)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_occurrences_file_hash ON document_occurrences(file_hash)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_occurrences_source_rel_path ON document_occurrences(source_rel_path)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_occurrences_dataset_source_id ON document_occurrences(dataset_source_id)") connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_document_occurrences_active_rel_path ON document_occurrences(rel_path) WHERE lifecycle_status = 'active'") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_document_occurrences_active_source_identity ON document_occurrences(source_kind, COALESCE(custodian, ''), COALESCE(source_rel_path, ''), COALESCE(source_item_id, '')) WHERE lifecycle_status = 'active' """ ) connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_document_dedupe_keys_basis_value ON document_dedupe_keys(basis, key_value)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_dedupe_keys_document_id ON document_dedupe_keys(document_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_canonical_metadata_conflicts_document_field ON canonical_metadata_conflicts(document_id, field_name)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_control_number_aliases_alias_value ON document_control_number_aliases(alias_value)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_merge_events_survivor_loser_created ON document_merge_events(survivor_document_id, loser_document_id, created_at)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_field_conflicts_merge_event_id ON document_field_conflicts(merge_event_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_dataset_sources_dataset_id ON dataset_sources(dataset_id)") connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_dataset_sources_locator_unique ON dataset_sources(source_kind, source_locator)") connection.execute("DROP INDEX IF EXISTS idx_conversations_source_kind_key") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_conversations_source_locator_key_unique ON conversations(source_kind, source_locator, conversation_key) """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_conversations_source_locator ON conversations(source_locator)") connection.execute("CREATE INDEX IF NOT EXISTS idx_dataset_documents_document_id ON dataset_documents(document_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_dataset_documents_dataset_id ON dataset_documents(dataset_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_dataset_documents_source_id ON dataset_documents(dataset_source_id)") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_dataset_documents_membership_unique ON dataset_documents(dataset_id, document_id, COALESCE(dataset_source_id, 0)) """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_productions_dataset_id ON productions(dataset_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_container_sources_dataset_id ON container_sources(dataset_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_container_sources_source_kind ON container_sources(source_kind)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entities_status_type ON entities(canonical_status, entity_type)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entities_merged_into ON entities(merged_into_entity_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entities_label ON entities(display_name, primary_email)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entity_identifiers_entity_type ON entity_identifiers(entity_id, identifier_type)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entity_identifiers_lookup ON entity_identifiers(identifier_type, normalized_value)") connection.execute( """ CREATE INDEX IF NOT EXISTS idx_entity_identifiers_scoped_handle ON entity_identifiers(provider, provider_scope, normalized_value) WHERE identifier_type = 'handle' """ ) connection.execute( """ CREATE INDEX IF NOT EXISTS idx_entity_identifiers_external_id ON entity_identifiers(identifier_name, identifier_scope, normalized_value) WHERE identifier_type = 'external_id' """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_entity_resolution_keys_entity ON entity_resolution_keys(entity_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entity_resolution_keys_lookup ON entity_resolution_keys(key_type, normalized_value)") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_entity_resolution_keys_email_unique ON entity_resolution_keys(key_type, normalized_value) WHERE key_type = 'email' """ ) connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_entity_resolution_keys_handle_unique ON entity_resolution_keys(key_type, provider, provider_scope, normalized_value) WHERE key_type = 'handle' """ ) connection.execute("DROP INDEX IF EXISTS idx_entity_resolution_keys_external_id_unique") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_entity_resolution_keys_external_id_unique ON entity_resolution_keys(key_type, identifier_name, COALESCE(identifier_scope, ''), normalized_value) WHERE key_type = 'external_id' """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_document_entities_document_role ON document_entities(document_id, role, ordinal)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_entities_entity_role ON document_entities(entity_id, role)") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_document_entities_unique ON document_entities(document_id, role, entity_id) """ ) connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_entity_merge_blocks_pair ON entity_merge_blocks(left_entity_id, right_entity_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entity_overrides_source_entity ON entity_overrides(source_entity_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_entity_overrides_replacement_entity ON entity_overrides(replacement_entity_id)") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_entity_overrides_document_effect ON entity_overrides( scope_type, COALESCE(scope_id, 0), COALESCE(role, ''), COALESCE(source_entity_id, 0), COALESCE(normalized_candidate_key, ''), override_effect ) WHERE scope_type = 'document' """ ) connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_entity_overrides_global_ignore ON entity_overrides( scope_type, COALESCE(role, ''), COALESCE(source_entity_id, 0), COALESCE(normalized_candidate_key, ''), COALESCE(source_hint, ''), override_effect ) WHERE scope_type = 'global' AND override_effect = 'ignore' """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_document_email_threading_message_id ON document_email_threading(message_id)") connection.execute( "CREATE INDEX IF NOT EXISTS idx_document_email_threading_conversation_index ON document_email_threading(conversation_index)" ) connection.execute( "CREATE INDEX IF NOT EXISTS idx_document_email_threading_conversation_topic ON document_email_threading(conversation_topic)" ) connection.execute( "CREATE INDEX IF NOT EXISTS idx_document_email_threading_normalized_subject ON document_email_threading(normalized_subject)" ) connection.execute("CREATE INDEX IF NOT EXISTS idx_document_chat_threading_thread_id ON document_chat_threading(thread_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_document_chat_threading_message_id ON document_chat_threading(message_id)") connection.execute( "CREATE INDEX IF NOT EXISTS idx_document_chat_threading_parent_message_id ON document_chat_threading(parent_message_id)" ) connection.execute("DROP INDEX IF EXISTS idx_documents_display_sort") connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_control_number_unique ON documents(control_number)") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_source_identity_unique ON documents(source_rel_path, source_item_id) WHERE source_kind IS NOT NULL AND source_item_id IS NOT NULL """ ) connection.execute( """ CREATE INDEX IF NOT EXISTS idx_documents_control_number_sort ON documents(control_number_batch, control_number_family_sequence, control_number_attachment_sequence) """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_job_outputs_job_id ON job_outputs(job_id, ordinal)") connection.execute("CREATE INDEX IF NOT EXISTS idx_job_versions_job_id ON job_versions(job_id, version DESC)") connection.execute("CREATE INDEX IF NOT EXISTS idx_job_versions_capability ON job_versions(capability)") connection.execute("CREATE INDEX IF NOT EXISTS idx_runs_job_version_id ON runs(job_version_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_runs_from_run_id ON runs(from_run_id)") connection.execute( """ CREATE INDEX IF NOT EXISTS idx_run_snapshot_documents_run_id_ordinal ON run_snapshot_documents(run_id, ordinal, id) """ ) connection.execute( "CREATE INDEX IF NOT EXISTS idx_run_snapshot_documents_document_id ON run_snapshot_documents(document_id)" ) connection.execute("CREATE INDEX IF NOT EXISTS idx_run_items_run_id_status ON run_items(run_id, status)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_items_document_id ON run_items(document_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_items_result_id ON run_items(result_id)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_items_page_number ON run_items(run_id, document_id, page_number)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_items_run_claim ON run_items(run_id, status, claimed_by)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_items_heartbeat ON run_items(last_heartbeat_at)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_workers_run_id ON run_workers(run_id, status)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_workers_claimed_by ON run_workers(run_id, claimed_by)") connection.execute("CREATE INDEX IF NOT EXISTS idx_run_workers_task_id ON run_workers(worker_task_id)") connection.execute("DROP INDEX IF EXISTS idx_run_items_snapshot_kind_unique") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_run_items_snapshot_kind_unique ON run_items( run_id, COALESCE(run_snapshot_document_id, 0), item_kind, COALESCE(page_number, 0), COALESCE(segment_id, 0) ) """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_attempts_run_item_id ON attempts(run_item_id, attempt_number)") connection.execute("CREATE INDEX IF NOT EXISTS idx_ocr_page_outputs_run_id_doc_page ON ocr_page_outputs(run_id, document_id, page_number)") connection.execute( "CREATE INDEX IF NOT EXISTS idx_image_description_page_outputs_run_id_doc_page " "ON image_description_page_outputs(run_id, document_id, page_number)" ) connection.execute("CREATE INDEX IF NOT EXISTS idx_results_job_version_id ON results(job_version_id, created_at)") connection.execute("CREATE INDEX IF NOT EXISTS idx_results_document_id ON results(document_id, created_at)") connection.execute("CREATE INDEX IF NOT EXISTS idx_results_input_revision_id ON results(input_revision_id)") connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_results_active_identity_unique ON results(document_id, job_version_id, input_identity) WHERE retracted_at IS NULL """ ) connection.execute("CREATE INDEX IF NOT EXISTS idx_result_outputs_result_id ON result_outputs(result_id)") connection.execute( """ CREATE INDEX IF NOT EXISTS idx_text_revisions_document_created_at ON text_revisions(document_id, created_at, id) """ ) connection.execute( """ CREATE INDEX IF NOT EXISTS idx_text_revision_segments_revision_profile ON text_revision_segments(revision_id, segment_profile, level, ordinal) """ ) connection.execute( """ CREATE UNIQUE INDEX IF NOT EXISTS idx_embedding_vectors_active_segment_unique ON embedding_vectors(job_version_id, segment_id) WHERE retracted_at IS NULL """ ) connection.execute( """ CREATE INDEX IF NOT EXISTS idx_publications_document_field ON publications(document_id, custom_field_name, published_at) """ ) connection.execute( """ CREATE INDEX IF NOT EXISTS idx_text_revision_activation_events_document_created_at ON text_revision_activation_events(document_id, created_at, id) """ ) merged_legacy_locks = merge_legacy_field_locks(connection) backfilled_content_type = backfill_content_type(connection) backfilled_child_document_kinds = backfill_child_document_kinds(connection) backfilled_conversation_assignment_modes = backfill_conversation_assignment_modes(connection) backfilled_source_kinds = backfill_source_kinds(connection) rewrote_internal_rel_path_prefix = backfill_internal_rel_path_prefix(connection) dataset_membership_migration_needed = prior_schema_version is None or prior_schema_version < 12 backfilled_dataset_ids = backfill_dataset_ids(connection, root) if dataset_membership_migration_needed else 0 backfilled_dataset_memberships = backfill_dataset_memberships(connection) if dataset_membership_migration_needed else 0 backfilled_occurrence_dataset_source_ids = backfill_occurrence_dataset_source_ids(connection) backfilled_dataset_name_normalized = backfill_dataset_name_normalized(connection) merged_duplicate_dataset_identities = merge_dataset_identity_duplicates(connection) suffixed_dataset_name_collisions = suffix_dataset_name_collisions(connection) connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_datasets_source_locator_unique ON datasets(source_kind, dataset_locator)") connection.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_datasets_name_normalized_unique ON datasets(dataset_name_normalized)") backfilled_custodian = backfill_custodian(connection) backfilled_control_numbers = backfill_control_numbers(connection) rebuilt_control_number_batches = backfill_control_number_batches(connection) backfilled_document_occurrences = backfill_document_occurrences(connection) backfilled_document_dedupe_keys = backfill_document_dedupe_keys(connection) refresh_document_occurrence_caches = should_refresh_document_occurrence_caches( prior_schema_version=prior_schema_version, backfilled_occurrence_dataset_source_ids=backfilled_occurrence_dataset_source_ids, backfilled_document_occurrences=backfilled_document_occurrences, ) refreshed_document_occurrence_caches = ( refresh_documents_from_occurrences(connection) if refresh_document_occurrence_caches else 0 ) removed_custodian_locks = drop_document_field_locks(connection, "custodian") removed_documents_custodian_column = remove_documents_custodian_column(connection) rebuilt_documents_fts = ensure_documents_fts(connection) connection.commit() return { "schema_version": SCHEMA_VERSION, "backfilled_content_type": backfilled_content_type, "backfilled_child_document_kinds": backfilled_child_document_kinds, "backfilled_conversation_assignment_modes": backfilled_conversation_assignment_modes, "backfilled_custodian": backfilled_custodian, "backfilled_dataset_ids": backfilled_dataset_ids, "backfilled_dataset_memberships": backfilled_dataset_memberships, "backfilled_occurrence_dataset_source_ids": backfilled_occurrence_dataset_source_ids, "backfilled_dataset_name_normalized": backfilled_dataset_name_normalized, "merged_duplicate_dataset_identities": merged_duplicate_dataset_identities, "suffixed_dataset_name_collisions": suffixed_dataset_name_collisions, "backfilled_source_kinds": backfilled_source_kinds, "rewrote_internal_rel_path_prefix": rewrote_internal_rel_path_prefix, "backfilled_control_numbers": backfilled_control_numbers, "rebuilt_control_number_batches": rebuilt_control_number_batches, "backfilled_document_occurrences": backfilled_document_occurrences, "migrated_export_work_items_nullable_document_id": migrated_export_work_items_nullable_document_id, "backfilled_document_dedupe_keys": backfilled_document_dedupe_keys, "refresh_document_occurrence_caches": refresh_document_occurrence_caches, "refreshed_document_occurrence_caches": refreshed_document_occurrence_caches, "removed_custodian_locks": removed_custodian_locks, "removed_documents_custodian_column": removed_documents_custodian_column, "backfilled_legacy_control_number": backfilled_legacy_control_number, "backfilled_legacy_control_number_batch": backfilled_legacy_control_number_batch, "backfilled_legacy_control_number_family_sequence": backfilled_legacy_control_number_family_sequence, "backfilled_legacy_control_number_attachment_sequence": backfilled_legacy_control_number_attachment_sequence, "rebuilt_documents_fts": rebuilt_documents_fts, "merged_legacy_locks": merged_legacy_locks, } def read_runtime(path: Path) -> dict[str, object] | None: if not path.exists(): return None return json.loads(path.read_text(encoding="utf-8")) def read_plugin_runtime_requirements_version(paths: dict[str, Path] | None) -> str | None: if paths is None: return None marker_path = paths["requirements_marker_path"] if not marker_path.exists(): return None return normalize_whitespace(marker_path.read_text(encoding="utf-8")) or None def write_plugin_runtime_requirements_version(paths: dict[str, Path], version: str) -> None: paths["requirements_marker_path"].write_text(f"{version}\n", encoding="utf-8") def activate_plugin_site_packages(paths: dict[str, Path] | None) -> bool: if paths is None: return False site_packages_path = first_existing_plugin_runtime_site_packages(paths) if site_packages_path is None: return False try: resolved = str(site_packages_path.resolve()) except OSError: resolved = str(site_packages_path) if resolved in ACTIVATED_PLUGIN_SITE_PACKAGES or resolved in sys.path: ACTIVATED_PLUGIN_SITE_PACKAGES.add(resolved) return True site.addsitedir(resolved) ACTIVATED_PLUGIN_SITE_PACKAGES.add(resolved) return True def describe_plugin_runtime(paths: dict[str, Path] | None) -> dict[str, object]: if paths is None: return { "status": "unavailable", "detail": "Shared plugin runtime path could not be determined from the current plugin installation.", "mode": "shared", "plugin_root": None, "runtime_root": None, "venv_present": False, "python_executable": None, "site_packages_path": None, "requirements_version": None, } venv_dir_present = paths["venv_dir"].exists() python_present = paths["venv_python_path"].exists() site_packages_path = first_existing_plugin_runtime_site_packages(paths) installed_requirements_version = read_plugin_runtime_requirements_version(paths) if python_present and installed_requirements_version == REQUIREMENTS_VERSION: status = "pass" detail = "Shared plugin runtime and pinned requirements are ready." elif python_present: status = "partial" detail = "Shared plugin runtime exists, but the pinned requirements are not fully installed yet." elif venv_dir_present: status = "partial" detail = "Shared plugin runtime directory exists, but the Python executable is missing." else: status = "missing" detail = "Shared plugin runtime has not been initialized yet." return { "status": status, "detail": detail, "mode": "shared", "plugin_root": str(paths["plugin_root"]), "runtime_root": str(paths["runtime_root"]), "venv_present": venv_dir_present, "python_executable": str(paths["venv_python_path"]) if python_present else None, "site_packages_path": str(site_packages_path) if site_packages_path is not None else None, "requirements_version": installed_requirements_version, } def ensure_plugin_runtime_layout(paths: dict[str, Path]) -> None: paths["runtime_root"].mkdir(parents=True, exist_ok=True) paths["locks_dir"].mkdir(parents=True, exist_ok=True) def ensure_plugin_venv(paths: dict[str, Path]) -> dict[str, object]: ensure_plugin_runtime_layout(paths) venv_dir = paths["venv_dir"] created = not paths["venv_python_path"].exists() if created: builder = venv.EnvBuilder(with_pip=True) builder.create(venv_dir) if not paths["venv_python_path"].exists(): raise RetrieverError( f"Plugin runtime initialization did not create a Python executable at {paths['venv_python_path']}" ) return { "created": created, "python_executable": str(paths["venv_python_path"]), } def install_plugin_runtime_requirements( paths: dict[str, Path], requirements: list[str] | tuple[str, ...], *, reason: str, ) -> dict[str, object]: python_path = paths["venv_python_path"] if not python_path.exists(): raise RetrieverError(f"Plugin runtime Python not found at {python_path}") command = [ str(python_path), "-m", "pip", "install", "--disable-pip-version-check", *list(requirements), ] completed = subprocess.run(command, capture_output=True, text=True) if completed.returncode != 0: detail = normalize_whitespace(completed.stderr or completed.stdout) or f"pip exited with status {completed.returncode}" raise RetrieverError(f"Plugin runtime install failed for {reason}: {detail}") write_plugin_runtime_requirements_version(paths, REQUIREMENTS_VERSION) return { "installed": True, "reason": reason, "command": command, "requirements_version": REQUIREMENTS_VERSION, } def ensure_plugin_runtime( paths: dict[str, Path] | None, *, install_requirements: bool, force_requirements_install: bool = False, reason: str, ) -> dict[str, object]: if paths is None: raise RetrieverError("Could not determine the shared plugin runtime location for this Retriever installation.") ensure_plugin_runtime_layout(paths) lock_handle = acquire_plugin_runtime_install_lock(paths) try: venv_result = ensure_plugin_venv(paths) activate_plugin_site_packages(paths) installed_requirements_version = read_plugin_runtime_requirements_version(paths) requirements_installed = False if install_requirements and ( force_requirements_install or installed_requirements_version != REQUIREMENTS_VERSION ): install_plugin_runtime_requirements(paths, PINNED_RUNTIME_REQUIREMENTS, reason=reason) installed_requirements_version = REQUIREMENTS_VERSION requirements_installed = True activate_plugin_site_packages(paths) runtime_status = describe_plugin_runtime(paths) return { "status": runtime_status["status"], "detail": runtime_status["detail"], "venv_created": bool(venv_result["created"]), "requirements_installed": requirements_installed, "requirements_version": installed_requirements_version, "python_executable": runtime_status["python_executable"], } finally: release_plugin_runtime_install_lock(lock_handle) def read_workspace_meta(connection: sqlite3.Connection) -> dict[str, object] | None: if not table_exists(connection, "workspace_meta"): return None row = connection.execute( """ SELECT id, schema_version, tool_version, requirements_version, template_source, template_sha256, created_at, updated_at FROM workspace_meta WHERE id = 1 """ ).fetchone() if row is None: return None return {key: row[key] for key in row.keys()} def write_runtime(paths: dict[str, Path], tool_sha256: str | None) -> dict[str, object]: now = utc_now() plugin_runtime_paths_for_workspace = plugin_runtime_paths(root=paths["root"]) runtime = { "tool_version": TOOL_VERSION, "schema_version": SCHEMA_VERSION, "requirements_version": REQUIREMENTS_VERSION, "template_source": TEMPLATE_SOURCE, "template_sha256": tool_sha256, "python_version": platform.python_version(), "generated_at": now, "last_verified_at": now, "plugin_runtime": describe_plugin_runtime(plugin_runtime_paths_for_workspace), } paths["runtime_path"].write_text(json.dumps(runtime, indent=2, sort_keys=True), encoding="utf-8") return runtime def probe_fts5() -> dict[str, str]: connection = sqlite3.connect(":memory:") try: connection.execute("CREATE VIRTUAL TABLE t USING fts5(x)") return {"status": "pass", "detail": "FTS5 virtual table created successfully"} except Exception as exc: # pragma: no cover - runtime probe return {"status": "fail", "detail": f"{type(exc).__name__}: {exc}"} finally: connection.close() def detect_platform() -> dict[str, str]: return { "platform": platform.platform(), "system": platform.system(), "release": platform.release(), "machine": platform.machine(), } def probe_processing_providers(connection: sqlite3.Connection | None) -> dict[str, object]: configured_providers: list[str] = [] configured_capabilities: list[str] = [] if connection is not None and table_exists(connection, "job_versions"): configured_providers = sorted( { normalize_whitespace(str(row["provider"] or "")).lower() for row in connection.execute( """ SELECT DISTINCT provider FROM job_versions WHERE provider IS NOT NULL AND TRIM(provider) != '' """ ).fetchall() if normalize_whitespace(str(row["provider"] or "")) } ) configured_capabilities = sorted( { normalize_whitespace(str(row["capability"] or "")).lower() for row in connection.execute( """ SELECT DISTINCT capability FROM job_versions WHERE capability IS NOT NULL AND TRIM(capability) != '' """ ).fetchall() if normalize_whitespace(str(row["capability"] or "")) } ) return { "configured_providers": configured_providers, "configured_capabilities": configured_capabilities, "cowork_runtime": { "status": "pass", "detail": "Cowork agent execution is the primary runtime path for processing jobs.", }, "external_providers": { "status": "warn" if configured_providers else "pass", "detail": ( "External provider identifiers are stored for future integrations, but are not required for Cowork execution." if configured_providers else "No external provider identifiers are configured." ), }, } def determine_workspace_state(paths: dict[str, Path]) -> str: state_dir = paths["state_dir"].exists() db_path = paths["db_path"].exists() db_usable = db_path and (file_size_bytes(paths["db_path"]) or 0) > 0 runtime_path = paths["runtime_path"].exists() if all((state_dir, db_usable, runtime_path)): return "initialized" if any((state_dir, db_path, runtime_path)): return "partial" return "missing" def workspace_status(root: Path, quick: bool) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) runtime_paths = plugin_runtime_paths(root=root) plugin_runtime = describe_plugin_runtime(runtime_paths) fts5 = probe_fts5() pip_python = ( runtime_paths["venv_python_path"] if runtime_paths is not None and runtime_paths["venv_python_path"].exists() else Path(sys.executable) ) pip_ok, pip_version = run_command([str(pip_python), "-m", "pip", "--version"]) pst_backend = dependency_status( "pypff", package_name="libpff-python", import_name="pypff", detail_label="PST backend", probe_if_unloaded=True, allow_auto_install=False, ) runtime = read_runtime(paths["runtime_path"]) canonical_tool_path = locate_canonical_plugin_tool_or_self() current_sha = sha256_file(canonical_tool_path) if canonical_tool_path is not None else None stored_sha = None if runtime is None else runtime.get("template_sha256") workspace_state = determine_workspace_state(paths) registry_status = None workspace_inventory = None processing_providers = probe_processing_providers(None) journal_mode = None db_error = None workspace_meta = None actual_schema_version = None schema_needs_migration = None workspace_inventory_error = None registry_error = None if paths["db_path"].exists() and (file_size_bytes(paths["db_path"]) or 0) > 0: try: connection = connect_db(paths["db_path"]) try: journal_mode = current_journal_mode(connection) workspace_meta = read_workspace_meta(connection) if workspace_meta is not None and workspace_meta.get("schema_version") is not None: actual_schema_version = int(workspace_meta["schema_version"]) schema_needs_migration = actual_schema_version < SCHEMA_VERSION elif workspace_state != "missing": schema_needs_migration = True if schema_needs_migration is not True: if table_exists(connection, "documents"): try: workspace_inventory = document_inventory_counts(connection) except Exception as exc: workspace_inventory_error = f"{type(exc).__name__}: {exc}" if table_exists(connection, "documents") and table_exists(connection, "custom_fields_registry"): try: registry_status = reconcile_custom_fields_registry(connection, repair=False) except Exception as exc: registry_error = f"{type(exc).__name__}: {exc}" processing_providers = probe_processing_providers(connection) finally: connection.close() except Exception as exc: db_error = f"{type(exc).__name__}: {exc}" overall = "pass" if fts5["status"] != "pass": overall = "fail" elif db_error is not None: overall = "fail" elif ( workspace_state == "partial" or schema_needs_migration is True or workspace_inventory_error is not None or registry_error is not None ): overall = "partial" result: dict[str, object] = { "overall": overall, "tool_version": TOOL_VERSION, "schema_version": SCHEMA_VERSION, "workspace_schema_version": actual_schema_version, "schema_needs_migration": schema_needs_migration, "python_version": platform.python_version(), "pip_version": pip_version if pip_ok else None, "pip_status": "pass" if pip_ok else "fail", "sqlite_version": sqlite3.sqlite_version, "fts5": fts5, "pst_backend": pst_backend, "processing_providers": processing_providers, "platform": detect_platform(), "plugin_runtime": plugin_runtime, "workspace": { "root": str(root.resolve()), "state": workspace_state, "db_present": paths["db_path"].exists(), "db_size_bytes": file_size_bytes(paths["db_path"]), "runtime_present": paths["runtime_path"].exists(), "canonical_tool_present": canonical_tool_path is not None, }, "tool_integrity": { "current_sha256": current_sha, "runtime_sha256": stored_sha, "matches_runtime": current_sha == stored_sha if current_sha and stored_sha else None, }, } if journal_mode is not None: result["sqlite_journal_mode"] = journal_mode if db_error is not None: result["db_error"] = db_error if workspace_inventory is not None: result["workspace_inventory"] = workspace_inventory if workspace_inventory_error is not None: result["workspace_inventory_error"] = workspace_inventory_error if registry_error is not None: result["custom_field_registry_error"] = registry_error if not quick: result["paths"] = {key: str(value) for key, value in paths.items()} if runtime is not None: result["runtime"] = runtime if workspace_meta is not None: result["workspace_meta"] = workspace_meta if registry_status is not None: result["custom_field_registry"] = registry_status return result def doctor(root: Path, quick: bool, *, repair_stale_sidecars: bool = False) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) runtime_paths = plugin_runtime_paths(root=root) plugin_runtime = describe_plugin_runtime(runtime_paths) fts5 = probe_fts5() pip_python = ( runtime_paths["venv_python_path"] if runtime_paths is not None and runtime_paths["venv_python_path"].exists() else Path(sys.executable) ) pip_ok, pip_version = run_command([str(pip_python), "-m", "pip", "--version"]) pst_backend = dependency_status( "pypff", package_name="libpff-python", import_name="pypff", detail_label="PST backend", probe_if_unloaded=True, allow_auto_install=False, ) runtime = read_runtime(paths["runtime_path"]) canonical_tool_path = locate_canonical_plugin_tool_or_self() current_sha = sha256_file(canonical_tool_path) if canonical_tool_path is not None else None stored_sha = None if runtime is None else runtime.get("template_sha256") workspace_state = determine_workspace_state(paths) registry_status = None schema_status = None workspace_inventory = None processing_providers = probe_processing_providers(None) journal_mode = None db_error = None integrity_check: dict[str, object] | None = None stale_artifacts_before = [str(path) for path in stale_sqlite_artifact_paths(paths["db_path"])] removed_stale_artifacts: list[str] = [] if repair_stale_sidecars and stale_artifacts_before: removed_stale_artifacts = remove_stale_sqlite_artifacts(paths["db_path"]) if paths["db_path"].exists(): try: connection = connect_db(paths["db_path"]) try: journal_mode = current_journal_mode(connection) try: integrity_rows = connection.execute("PRAGMA integrity_check").fetchall() integrity_messages = [str(row[0]) for row in integrity_rows] integrity_check = { "status": "ok" if integrity_messages == ["ok"] else "error", "messages": integrity_messages, } except Exception as exc: integrity_check = { "status": "failed", "messages": [f"{type(exc).__name__}: {exc}"], } if integrity_check is None or str(integrity_check.get("status")) == "ok": schema_status = apply_schema(connection, root) registry_status = reconcile_custom_fields_registry(connection, repair=True) workspace_inventory = document_inventory_counts(connection) processing_providers = probe_processing_providers(connection) finally: connection.close() except Exception as exc: db_error = f"{type(exc).__name__}: {exc}" overall = "pass" if fts5["status"] != "pass": overall = "fail" elif db_error is not None: overall = "fail" elif integrity_check is not None and str(integrity_check.get("status") or "") != "ok": overall = "fail" elif workspace_state == "partial": overall = "partial" elif stale_artifacts_before and not removed_stale_artifacts: overall = "partial" result: dict[str, object] = { "overall": overall, "tool_version": TOOL_VERSION, "schema_version": SCHEMA_VERSION, "python_version": platform.python_version(), "pip_version": pip_version if pip_ok else None, "pip_status": "pass" if pip_ok else "fail", "sqlite_version": sqlite3.sqlite_version, "fts5": fts5, "pst_backend": pst_backend, "processing_providers": processing_providers, "platform": detect_platform(), "plugin_runtime": plugin_runtime, "workspace": { "root": str(root.resolve()), "state": workspace_state, "db_present": paths["db_path"].exists(), "db_size_bytes": file_size_bytes(paths["db_path"]), "runtime_present": paths["runtime_path"].exists(), "canonical_tool_present": canonical_tool_path is not None, }, "tool_integrity": { "current_sha256": current_sha, "runtime_sha256": stored_sha, "matches_runtime": current_sha == stored_sha if current_sha and stored_sha else None, }, "db": { "path": str(paths["db_path"]), "stale_sqlite_artifacts_before": stale_artifacts_before, "removed_stale_sqlite_artifacts": removed_stale_artifacts, "stale_sqlite_artifacts_after": [str(path) for path in stale_sqlite_artifact_paths(paths["db_path"])], }, } if journal_mode is not None: result["sqlite_journal_mode"] = journal_mode if db_error is not None: result["db_error"] = db_error if workspace_inventory is not None: result["workspace_inventory"] = workspace_inventory if integrity_check is not None: result["db"]["integrity_check"] = integrity_check if not quick: result["paths"] = {key: str(value) for key, value in paths.items()} if runtime is not None: result["runtime"] = runtime if schema_status is not None: result["schema_apply"] = schema_status if registry_status is not None: result["custom_field_registry"] = registry_status return result def bootstrap(root: Path) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) ensure_layout(paths) canonical_tool_path = locate_canonical_plugin_tool_or_self() tool_sha = ( sha256_file(canonical_tool_path) if canonical_tool_path is not None else sha256_file(Path(__file__).resolve()) ) recovered_sqlite_artifacts = remove_stale_sqlite_artifacts(paths["db_path"]) seeded_sqlite_db: dict[str, object] | None = None last_error: Exception | None = None for attempt in range(4): try: connection = connect_db(paths["db_path"]) try: journal_mode = current_journal_mode(connection) apply_schema(connection, root) registry_status = reconcile_custom_fields_registry(connection, repair=True) write_workspace_meta(connection, tool_sha) finally: connection.close() write_runtime(paths, tool_sha) result = { "status": "initialized" if paths["runtime_path"].exists() else "failed", "workspace_root": str(root.resolve()), "schema_version": SCHEMA_VERSION, "tool_version": TOOL_VERSION, "requirements_version": REQUIREMENTS_VERSION, "custom_field_registry": registry_status, } if journal_mode is not None: result["journal_mode"] = journal_mode if recovered_sqlite_artifacts: result["recovered_sqlite_artifacts"] = recovered_sqlite_artifacts if seeded_sqlite_db is not None: result["seeded_sqlite_db"] = seeded_sqlite_db return result except Exception as exc: last_error = exc retry_artifacts = remove_stale_sqlite_artifacts(paths["db_path"]) if retry_artifacts: for artifact in retry_artifacts: if artifact not in recovered_sqlite_artifacts: recovered_sqlite_artifacts.append(artifact) continue if seeded_sqlite_db is None and sqlite_bootstrap_seed_required(paths["db_path"], exc): seeded_sqlite_db = seed_sqlite_db_from_local_temp(paths["db_path"]) for artifact in list(seeded_sqlite_db.get("reset_artifacts") or []): if artifact not in recovered_sqlite_artifacts: recovered_sqlite_artifacts.append(str(artifact)) continue break detail = f"{type(last_error).__name__}: {last_error}" if last_error is not None else "unknown bootstrap failure" if recovered_sqlite_artifacts: detail = ( f"{detail}. Removed stale SQLite artifacts before retry: " f"{', '.join(recovered_sqlite_artifacts)}" ) if seeded_sqlite_db is not None: detail = ( f"{detail}. Seeded SQLite DB from local temp with journal mode " f"{seeded_sqlite_db.get('journal_mode') or 'delete'} before retry" ) raise RetrieverError(f"Bootstrap failed for {paths['db_path']}: {detail}") from last_error # Commands that must not trigger an auto-upgrade. `schema-version` needs to # work even when a workspace does not exist yet. `workspace` owns the # explicit init/status/update flows, so it should not transparently rewrite # itself before dispatch. `slash` is intentionally not exempt so user-facing # slash commands also benefit from clean-but-stale auto-upgrades. AUTO_UPGRADE_EXEMPT_COMMANDS = frozenset( { "schema-version", "workspace", } ) def locate_canonical_plugin_tool(current_file: str | None = None) -> Path | None: """Find the canonical ``skills/tool-template/tools.py`` bundle.""" def _safe_resolve(path: Path) -> Path | None: try: return path.resolve() except OSError: return None env_path = os.environ.get("RETRIEVER_CANONICAL_TOOL_PATH") if env_path: candidate = Path(env_path).expanduser() if candidate.is_file(): resolved = _safe_resolve(candidate) if resolved is not None: return resolved candidate_path = current_file or __file__ try: start = Path(candidate_path).resolve() except OSError: start = None if start is None: return None sibling_candidate = start.with_name("tools.py") sibling_resolved = _safe_resolve(sibling_candidate) if sibling_resolved is not None and sibling_resolved.is_file(): return sibling_resolved for parent in [start.parent, *start.parents]: candidate = parent / "skills" / "tool-template" / "tools.py" resolved = _safe_resolve(candidate) if resolved is not None and resolved.is_file(): return resolved return None def locate_canonical_plugin_tool_or_self(current_file: str | None = None) -> Path | None: located = locate_canonical_plugin_tool(current_file) if located is not None: return located candidate_path = current_file or __file__ try: candidate = Path(candidate_path).resolve() except OSError: return None if ( candidate.is_file() and candidate.parent.name == "tool-template" and candidate.parent.parent.name == "skills" and candidate.name == "tools.py" ): return candidate return None if ( candidate.is_file() and candidate.parent.name == "tool-template" and candidate.parent.parent.name == "skills" ): return candidate return None def upgrade_workspace_tool( root: Path, canonical_path: Path, *, force: bool = False, reason: str = "manual", ) -> dict[str, object]: """Refresh workspace runtime metadata from the canonical tool bundle.""" del force paths = workspace_paths(root) ensure_layout(paths) runtime = read_runtime(paths["runtime_path"]) runtime_sha = runtime.get("template_sha256") if isinstance(runtime, dict) else None runtime_version = runtime.get("tool_version") if isinstance(runtime, dict) else None canonical_sha = sha256_file(canonical_path) if canonical_sha is None: raise RetrieverError(f"Canonical tool not readable at {canonical_path}") write_runtime(paths, canonical_sha) meta_updated = False meta_error: str | None = None try: connection = connect_db(paths["db_path"]) try: write_workspace_meta(connection, canonical_sha) meta_updated = True finally: connection.close() except Exception as exc: # pragma: no cover - best-effort path meta_error = f"{type(exc).__name__}: {exc}" result: dict[str, object] = { "status": "no-op" if runtime_sha == canonical_sha else "updated-runtime", "reason": reason, "previous_tool_sha256": runtime_sha, "previous_tool_version": runtime_version, "new_tool_sha256": canonical_sha, "new_tool_version": TOOL_VERSION, "canonical_path": str(canonical_path), "workspace_meta_updated": meta_updated, } if meta_error is not None: result["workspace_meta_error"] = meta_error return result def maybe_upgrade_workspace_tool(root: Path) -> dict[str, object] | None: """Best-effort runtime metadata refresh for a workspace.""" paths = workspace_paths(root) if not paths["state_dir"].exists(): return None runtime = read_runtime(paths["runtime_path"]) if not isinstance(runtime, dict): return None canonical_path = locate_canonical_plugin_tool_or_self() if canonical_path is None: return None canonical_sha = sha256_file(canonical_path) if canonical_sha is None: return None if runtime.get("template_sha256") == canonical_sha and runtime.get("template_source") == TEMPLATE_SOURCE: return None return upgrade_workspace_tool(root, canonical_path, reason="auto-runtime-sync") def init_workspace(root: Path, quick: bool = False) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) runtime_paths = plugin_runtime_paths(root=root) tool_update: dict[str, object] | None = None canonical_path = locate_canonical_plugin_tool_or_self() if canonical_path is not None: if not paths["runtime_path"].exists(): tool_update = upgrade_workspace_tool( root, canonical_path, force=False, reason="workspace-init", ) else: maybe_update = maybe_upgrade_workspace_tool(root) if maybe_update is not None: tool_update = maybe_update runtime_init = ensure_plugin_runtime( runtime_paths, install_requirements=True, force_requirements_install=False, reason="init", ) initialization_payload = bootstrap(root) result: dict[str, object] = { "action": "init", "status": "ready", "workspace_root": str(root.resolve()), "initialization": initialization_payload, "status_report": workspace_status(root, quick=quick), } if tool_update is not None: result["tool_update"] = tool_update if payload_has_meaningful_value(runtime_init): result["runtime_init"] = runtime_init return result def write_workspace_meta(connection: sqlite3.Connection, tool_sha256: str | None) -> None: now = utc_now() connection.execute( """ INSERT INTO workspace_meta ( id, schema_version, tool_version, requirements_version, template_source, template_sha256, created_at, updated_at ) VALUES (1, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET schema_version = excluded.schema_version, tool_version = excluded.tool_version, requirements_version = excluded.requirements_version, template_source = excluded.template_source, template_sha256 = excluded.template_sha256, updated_at = excluded.updated_at """, ( SCHEMA_VERSION, TOOL_VERSION, REQUIREMENTS_VERSION, TEMPLATE_SOURCE, tool_sha256, now, now, ), ) connection.commit() def workspace_runtime_metadata_needs_materialization( paths: dict[str, Path], connection: sqlite3.Connection, ) -> bool: if not paths["runtime_path"].exists(): return True workspace_meta = read_workspace_meta(connection) if workspace_meta is None: return True try: return int(workspace_meta.get("schema_version")) != SCHEMA_VERSION except (TypeError, ValueError): return True def ensure_workspace_runtime_metadata(root: Path, connection: sqlite3.Connection) -> None: paths = workspace_paths(root) if not workspace_runtime_metadata_needs_materialization(paths, connection): return canonical_tool_path = locate_canonical_plugin_tool_or_self() tool_path = canonical_tool_path or Path(__file__).resolve() tool_sha = sha256_file(tool_path) if tool_sha is None: raise RetrieverError(f"Canonical tool not readable at {tool_path}") write_workspace_meta(connection, tool_sha) write_runtime(paths, tool_sha) def resolve_production_root_argument(workspace_root: Path, raw_production_root: str | Path) -> Path: candidate = Path(raw_production_root).expanduser() if not candidate.is_absolute(): candidate = workspace_root / candidate candidate = candidate.resolve() try: candidate.relative_to(workspace_root.resolve()) except ValueError as exc: raise RetrieverError("Production root must be inside the workspace root for Phase 4 ingest.") from exc return candidate def production_row_signature( existing_row: sqlite3.Row | None, *, rel_path: str, file_name: str, source_kind: str, production_id: int, begin_bates: str, end_bates: str, begin_attachment: str | None, end_attachment: str | None, extracted: dict[str, object], source_parts: list[dict[str, object]], ) -> tuple[object, ...]: existing_locks = existing_row[MANUAL_FIELD_LOCKS_COLUMN] if existing_row is not None else "[]" return ( rel_path, file_name, source_kind, production_id, begin_bates, end_bates, begin_attachment, end_attachment, extracted.get("title"), extracted.get("content_type"), extracted.get("text_status"), sha256_text(str(extracted.get("text_content") or "")), extracted.get("page_count"), tuple( sorted( (part["part_kind"], part["rel_source_path"], int(part.get("ordinal", 0))) for part in source_parts ) ), bool(extracted.get("preview_artifacts")), existing_locks, ) def existing_production_row_signature(connection: sqlite3.Connection, row: sqlite3.Row | None) -> tuple[object, ...] | None: if row is None: return None source_parts = connection.execute( """ SELECT part_kind, rel_source_path, ordinal FROM document_source_parts WHERE document_id = ? ORDER BY part_kind ASC, ordinal ASC, id ASC """, (row["id"],), ).fetchall() return ( row["rel_path"], row["file_name"], row["source_kind"], row["production_id"], row["begin_bates"], row["end_bates"], row["begin_attachment"], row["end_attachment"], row["title"], row["content_type"], row["text_status"], row["content_hash"], row["page_count"], tuple((part["part_kind"], part["rel_source_path"], int(part["ordinal"])) for part in source_parts), connection.execute("SELECT COUNT(*) AS count FROM document_previews WHERE document_id = ?", (row["id"],)).fetchone()["count"] > 0, row[MANUAL_FIELD_LOCKS_COLUMN], ) JOB_KINDS = { "embedding", "image_description", "ocr", "structured_extraction", "translation", } REVISION_PRODUCING_JOB_KINDS = {"image_description", "ocr", "translation"} JOB_CAPABILITIES = { "text_structured", "text_translation", "vision_description", "vision_ocr", } JOB_INPUT_BASES = { "active_search_text", "source_extract", "source_file", "source_parts", "text_revision", } JOB_OUTPUT_VALUE_TYPES = { "boolean", "date", "integer", "json", "real", "text", } RUN_FAMILY_MODES = {"exact", "with_family"} RUN_ITEM_KINDS = {"document", "page", "segment"} RUN_ITEM_STATUSES = {"completed", "failed", "pending", "running", "skipped"} RUN_STATUSES = {"canceled", "completed", "failed", "planned", "running"} RUN_WORKER_MODES = {"background", "inline"} RUN_WORKER_STATUSES = {"active", "canceled", "completed", "failed", "orphaned", "stopped"} TEXT_REVISION_ACTIVATION_POLICIES = {"always", "if_empty", "if_poor", "manual"} RUN_ACTIVATION_POLICIES = {"always", "if_empty", "if_poor", "manual"} DEFAULT_RUN_ITEM_CLAIM_STALE_SECONDS = 900 DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS = 45 DEFAULT_RUN_ITEM_CONTEXT_INLINE_BYTES = 50 * 1024 DEFAULT_RUN_ITEM_CLAIM_BATCH_SIZE = 10 RUN_JOB_MIN_SECONDS_TO_CLAIM = 5 DEFAULT_WORKER_BATCH_SIZE = 5 DEFAULT_WORKER_INLINE_MAX_ITEMS = 5 DEFAULT_WORKER_INLINE_MAX_BATCHES = 12 DEFAULT_WORKER_BACKGROUND_MAX_BATCHES = 3 DEFAULT_WORKER_BACKGROUND_WAKE_INTERVAL_SECONDS = 60 DEFAULT_WORKER_BACKGROUND_MAX_PARALLEL = 4 DEFAULT_OCR_RENDER_RESOLUTION = 150 DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS = 35 MAX_RESUMABLE_STEP_BUDGET_SECONDS = 40 def normalize_resumable_step_budget(raw_budget_seconds: int | None, *, label: str = "budget-seconds") -> int: budget_seconds = DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS if raw_budget_seconds is None else int(raw_budget_seconds) if budget_seconds < 1: raise RetrieverError(f"{label} must be >= 1.") if budget_seconds > MAX_RESUMABLE_STEP_BUDGET_SECONDS: raise RetrieverError( f"{label} cannot exceed {MAX_RESUMABLE_STEP_BUDGET_SECONDS} seconds in bounded worker mode." ) return budget_seconds def default_run_item_claim_stale_seconds_for_launch_mode(launch_mode: str) -> int: normalized_launch_mode = normalize_run_worker_mode(launch_mode) if normalized_launch_mode == "background": return DEFAULT_RUN_ITEM_CLAIM_STALE_SECONDS return DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS def lease_expiration_after(seconds: int, *, now: datetime | None = None) -> str: return format_utc_timestamp((now or datetime.now(timezone.utc)) + timedelta(seconds=max(1, int(seconds)))) def lease_is_active(expires_at: object, *, now: datetime | None = None) -> bool: parsed = parse_utc_timestamp(expires_at) if parsed is None: return False return parsed > (now or datetime.now(timezone.utc)) def sanitize_processing_identifier(raw_name: str, *, label: str, prefix: str) -> str: sanitized = re.sub(r"[^a-zA-Z0-9_]+", "_", raw_name.strip()).strip("_").lower() if not sanitized: raise RetrieverError(f"{label} becomes empty after sanitization.") if sanitized[0].isdigit(): sanitized = f"{prefix}_{sanitized}" return sanitized def normalize_job_kind(job_kind: str) -> str: normalized = normalize_whitespace(job_kind).lower() if normalized not in JOB_KINDS: raise RetrieverError( f"Unsupported job kind: {job_kind!r}. Expected one of {', '.join(sorted(JOB_KINDS))}." ) return normalized def normalize_job_capability(capability: str) -> str: normalized = normalize_whitespace(capability).lower() if normalized not in JOB_CAPABILITIES: raise RetrieverError( f"Unsupported capability: {capability!r}. Expected one of {', '.join(sorted(JOB_CAPABILITIES))}." ) return normalized def default_job_capability_for_kind(job_kind: str) -> str: normalized_kind = normalize_job_kind(job_kind) if normalized_kind == "structured_extraction": return "text_structured" if normalized_kind == "translation": return "text_translation" if normalized_kind == "image_description": return "vision_description" if normalized_kind == "ocr": return "vision_ocr" raise RetrieverError( f"Job kind {normalized_kind!r} does not have a default Cowork capability. " "Pass an explicit capability when creating the job version." ) def default_job_input_basis_for_kind(job_kind: str) -> str: normalized_kind = normalize_job_kind(job_kind) if normalized_kind == "ocr": return "source_parts" if normalized_kind in {"structured_extraction", "translation", "embedding"}: return "active_search_text" if normalized_kind == "image_description": return "source_parts" raise RetrieverError( f"Job kind {normalized_kind!r} does not have a default input basis. " "Pass an explicit input basis when creating the job version." ) def normalize_job_input_basis(input_basis: str) -> str: normalized = normalize_whitespace(input_basis).lower() if normalized not in JOB_INPUT_BASES: raise RetrieverError( f"Unsupported input basis: {input_basis!r}. Expected one of {', '.join(sorted(JOB_INPUT_BASES))}." ) return normalized def normalize_job_output_value_type(value_type: str) -> str: normalized = normalize_whitespace(value_type).lower() if normalized not in JOB_OUTPUT_VALUE_TYPES: raise RetrieverError( f"Unsupported job output value type: {value_type!r}. " f"Expected one of {', '.join(sorted(JOB_OUTPUT_VALUE_TYPES))}." ) return normalized def normalize_run_family_mode(family_mode: str) -> str: normalized = normalize_whitespace(family_mode).lower() if normalized not in RUN_FAMILY_MODES: raise RetrieverError( f"Unsupported family mode: {family_mode!r}. Expected one of {', '.join(sorted(RUN_FAMILY_MODES))}." ) return normalized def normalize_run_item_kind(item_kind: str) -> str: normalized = normalize_whitespace(item_kind).lower() if normalized not in RUN_ITEM_KINDS: raise RetrieverError( f"Unsupported run item kind: {item_kind!r}. Expected one of {', '.join(sorted(RUN_ITEM_KINDS))}." ) return normalized def normalize_run_item_status(status: str) -> str: normalized = normalize_whitespace(status).lower() if normalized not in RUN_ITEM_STATUSES: raise RetrieverError( f"Unsupported run item status: {status!r}. Expected one of {', '.join(sorted(RUN_ITEM_STATUSES))}." ) return normalized def normalize_run_worker_mode(mode: str) -> str: normalized = normalize_whitespace(mode).lower() if normalized not in RUN_WORKER_MODES: raise RetrieverError( f"Unsupported run worker mode: {mode!r}. Expected one of {', '.join(sorted(RUN_WORKER_MODES))}." ) return normalized def normalize_run_worker_status(status: str) -> str: normalized = normalize_whitespace(status).lower() if normalized not in RUN_WORKER_STATUSES: raise RetrieverError( f"Unsupported run worker status: {status!r}. Expected one of {', '.join(sorted(RUN_WORKER_STATUSES))}." ) return normalized def normalize_text_revision_activation_policy(policy: str) -> str: normalized = normalize_whitespace(policy).lower() if normalized not in TEXT_REVISION_ACTIVATION_POLICIES: raise RetrieverError( f"Unsupported activation policy: {policy!r}. " f"Expected one of {', '.join(sorted(TEXT_REVISION_ACTIVATION_POLICIES))}." ) return normalized def normalize_run_activation_policy(policy: str) -> str: normalized = normalize_whitespace(policy).lower() if normalized not in RUN_ACTIVATION_POLICIES: raise RetrieverError( f"Unsupported run activation policy: {policy!r}. " f"Expected one of {', '.join(sorted(RUN_ACTIVATION_POLICIES))}." ) return normalized def compact_json_text(value: object) -> str: return json.dumps(value, ensure_ascii=True, sort_keys=True, separators=(",", ":")) def parse_json_argument(raw_value: str | None, *, label: str, default: object) -> object: if raw_value is None or raw_value.strip() == "": return default try: return json.loads(raw_value) except (TypeError, ValueError, json.JSONDecodeError) as exc: raise RetrieverError(f"{label} must be valid JSON.") from exc def parse_json_object_argument(raw_value: str | None, *, label: str, default: dict[str, object] | None = None) -> dict[str, object]: parsed = parse_json_argument(raw_value, label=label, default=default or {}) if not isinstance(parsed, dict): raise RetrieverError(f"{label} must decode to a JSON object.") return parsed def decode_json_text(raw_value: object, *, default: object = None) -> object: if raw_value in (None, ""): return default if isinstance(raw_value, (dict, list, int, float, bool)): return raw_value try: return json.loads(str(raw_value)) except (TypeError, ValueError, json.JSONDecodeError): return default def build_text_revision_input_identity(input_revision_id: int) -> str: return sha256_text(str(int(input_revision_id))) def build_ocr_input_identity( source_file_hash: str, *, rendering_settings: dict[str, object] | None = None, backend_id: str, ) -> str: normalized_settings = compact_json_text(rendering_settings or {}) return sha256_text(f"{source_file_hash}||{normalized_settings}||{backend_id}") def build_translation_input_identity(source_revision_id: int, *, target_language: str) -> str: normalized_language = normalize_whitespace(target_language).lower() return sha256_text(f"{int(source_revision_id)}||{normalized_language}") def build_image_source_input_identity( source_file_hash: str, *, image_prep_settings: dict[str, object] | None = None, backend_id: str, ) -> str: normalized_settings = compact_json_text(image_prep_settings or {}) return sha256_text(f"{source_file_hash}||{normalized_settings}||{backend_id}") def build_segment_input_identity(segment_id: int) -> str: return sha256_text(str(int(segment_id))) def job_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "job_name": row["job_name"], "job_kind": row["job_kind"], "description": row["description"], "created_at": row["created_at"], "updated_at": row["updated_at"], "archived_at": row["archived_at"], } def job_output_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "job_id": int(row["job_id"]), "output_name": row["output_name"], "value_type": row["value_type"], "bound_custom_field": row["bound_custom_field"], "description": row["description"], "ordinal": int(row["ordinal"]), "created_at": row["created_at"], "updated_at": row["updated_at"], } def job_version_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "job_id": int(row["job_id"]), "version": int(row["version"]), "display_name": row["display_name"], "instruction_text": row["instruction_text"], "instruction_hash": row["instruction_hash"], "response_schema": decode_json_text(row["response_schema_json"]), "capability": row["capability"], "provider": row["provider"], "model": row["model"], "parameters": decode_json_text(row["parameters_json"], default={}) or {}, "input_basis": row["input_basis"], "segment_profile": row["segment_profile"], "aggregation_strategy": row["aggregation_strategy"], "created_at": row["created_at"], "archived_at": row["archived_at"], } def find_job_row_by_name(connection: sqlite3.Connection, job_name: str) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM jobs WHERE job_name = ? """, (job_name,), ).fetchone() def require_job_row_by_name(connection: sqlite3.Connection, job_name: str) -> sqlite3.Row: row = find_job_row_by_name(connection, job_name) if row is None: raise RetrieverError(f"Unknown job: {job_name}") return row def job_outputs_for_job(connection: sqlite3.Connection, job_id: int) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM job_outputs WHERE job_id = ? ORDER BY ordinal ASC, output_name ASC, id ASC """, (job_id,), ).fetchall() return [job_output_row_to_payload(row) for row in rows] def job_versions_for_job(connection: sqlite3.Connection, job_id: int) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM job_versions WHERE job_id = ? ORDER BY version DESC, id DESC """, (job_id,), ).fetchall() return [job_version_row_to_payload(row) for row in rows] def latest_job_version_for_job(connection: sqlite3.Connection, job_id: int) -> dict[str, object] | None: row = connection.execute( """ SELECT * FROM job_versions WHERE job_id = ? ORDER BY version DESC, id DESC LIMIT 1 """, (job_id,), ).fetchone() return None if row is None else job_version_row_to_payload(row) def job_summary_by_id(connection: sqlite3.Connection, job_id: int) -> dict[str, object]: row = connection.execute( """ SELECT * FROM jobs WHERE id = ? """, (job_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown job id: {job_id}") version_count_row = connection.execute( """ SELECT COUNT(*) AS count FROM job_versions WHERE job_id = ? """, (job_id,), ).fetchone() return { **job_row_to_payload(row), "outputs": job_outputs_for_job(connection, job_id), "job_version_count": int(version_count_row["count"] or 0), "latest_job_version": latest_job_version_for_job(connection, job_id), } def list_job_summaries(connection: sqlite3.Connection) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM jobs ORDER BY job_name ASC, id ASC """ ).fetchall() return [job_summary_by_id(connection, int(row["id"])) for row in rows] def create_job_row( connection: sqlite3.Connection, *, job_name: str, job_kind: str, description: str | None, ) -> int: now = utc_now() cursor = connection.execute( """ INSERT INTO jobs (job_name, job_kind, description, created_at, updated_at, archived_at) VALUES (?, ?, ?, ?, ?, NULL) """, (job_name, job_kind, description, now, now), ) return int(cursor.lastrowid) def next_job_output_ordinal(connection: sqlite3.Connection, job_id: int) -> int: row = connection.execute( """ SELECT COALESCE(MAX(ordinal), -1) AS max_ordinal FROM job_outputs WHERE job_id = ? """, (job_id,), ).fetchone() return int(row["max_ordinal"] or -1) + 1 def upsert_job_output_row( connection: sqlite3.Connection, *, job_id: int, output_name: str, value_type: str, bound_custom_field: str | None, description: str | None, ) -> tuple[int, bool]: existing_row = connection.execute( """ SELECT * FROM job_outputs WHERE job_id = ? AND output_name = ? """, (job_id, output_name), ).fetchone() now = utc_now() if existing_row is not None: connection.execute( """ UPDATE job_outputs SET value_type = ?, bound_custom_field = ?, description = ?, updated_at = ? WHERE id = ? """, (value_type, bound_custom_field, description, now, existing_row["id"]), ) return int(existing_row["id"]), False cursor = connection.execute( """ INSERT INTO job_outputs ( job_id, output_name, value_type, bound_custom_field, description, ordinal, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( job_id, output_name, value_type, bound_custom_field, description, next_job_output_ordinal(connection, job_id), now, now, ), ) return int(cursor.lastrowid), True def next_job_version_number(connection: sqlite3.Connection, job_id: int) -> int: row = connection.execute( """ SELECT COALESCE(MAX(version), 0) AS max_version FROM job_versions WHERE job_id = ? """, (job_id,), ).fetchone() return int(row["max_version"] or 0) + 1 def create_job_version_row( connection: sqlite3.Connection, *, job_id: int, job_name: str, instruction_text: str, response_schema_json: str | None, capability: str, provider: str, model: str | None, parameters_json: str, input_basis: str, segment_profile: str | None, aggregation_strategy: str | None, display_name: str | None, ) -> int: version = next_job_version_number(connection, job_id) resolved_display_name = display_name or f"{job_name} v{version}" now = utc_now() cursor = connection.execute( """ INSERT INTO job_versions ( job_id, version, display_name, instruction_text, instruction_hash, response_schema_json, capability, provider, model, parameters_json, input_basis, segment_profile, aggregation_strategy, created_at, archived_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL) """, ( job_id, version, resolved_display_name, instruction_text, sha256_text(instruction_text), response_schema_json, capability, provider, model, parameters_json, input_basis, segment_profile, aggregation_strategy, now, ), ) return int(cursor.lastrowid) def job_version_summary_by_id(connection: sqlite3.Connection, job_version_id: int) -> dict[str, object]: row = connection.execute( """ SELECT * FROM job_versions WHERE id = ? """, (job_version_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown job version id: {job_version_id}") payload = job_version_row_to_payload(row) payload["job"] = job_summary_by_id(connection, int(row["job_id"])) return payload def require_job_version_row_by_id(connection: sqlite3.Connection, job_version_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM job_versions WHERE id = ? """, (job_version_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown job version id: {job_version_id}") return row def require_job_version_row( connection: sqlite3.Connection, *, job_version_id: int | None = None, job_name: str | None = None, version: int | None = None, ) -> sqlite3.Row: if job_version_id is not None: return require_job_version_row_by_id(connection, job_version_id) if not job_name: raise RetrieverError("A job version id or job name is required.") job_row = require_job_row_by_name(connection, job_name) if version is None: row = connection.execute( """ SELECT * FROM job_versions WHERE job_id = ? ORDER BY version DESC, id DESC LIMIT 1 """, (job_row["id"],), ).fetchone() else: row = connection.execute( """ SELECT * FROM job_versions WHERE job_id = ? AND version = ? """, (job_row["id"], version), ).fetchone() if row is None: if version is None: raise RetrieverError(f"Job {job_name!r} has no job versions yet.") raise RetrieverError(f"Unknown job version {version} for job {job_name!r}.") return row def text_revision_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "document_id": int(row["document_id"]), "revision_kind": row["revision_kind"], "language": row["language"], "parent_revision_id": row["parent_revision_id"], "created_by_job_version_id": row["created_by_job_version_id"], "storage_rel_path": row["storage_rel_path"], "content_hash": row["content_hash"], "char_count": row["char_count"], "token_estimate": row["token_estimate"], "quality_score": row["quality_score"], "provider_metadata": decode_json_text(row["provider_metadata_json"], default={}) or {}, "created_at": row["created_at"], "retracted_at": row["retracted_at"], "retraction_reason": row["retraction_reason"], } def require_text_revision_row_by_id(connection: sqlite3.Connection, text_revision_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM text_revisions WHERE id = ? """, (text_revision_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown text revision id: {text_revision_id}") return row def root_text_revision_id(connection: sqlite3.Connection, text_revision_id: int) -> int: current_id = int(text_revision_id) seen_ids: set[int] = set() while True: if current_id in seen_ids: raise RetrieverError(f"Detected a text revision parent cycle at revision {current_id}.") seen_ids.add(current_id) row = require_text_revision_row_by_id(connection, current_id) parent_revision_id = row["parent_revision_id"] if parent_revision_id is None: return current_id current_id = int(parent_revision_id) def text_revision_summary_by_id(connection: sqlite3.Connection, text_revision_id: int) -> dict[str, object]: row = require_text_revision_row_by_id(connection, text_revision_id) payload = text_revision_row_to_payload(row) document_row = connection.execute( """ SELECT source_text_revision_id, active_search_text_revision_id FROM documents WHERE id = ? """, (row["document_id"],), ).fetchone() payload["is_source_revision"] = ( document_row is not None and document_row["source_text_revision_id"] is not None and int(document_row["source_text_revision_id"]) == int(text_revision_id) ) payload["is_active_search_revision"] = ( document_row is not None and document_row["active_search_text_revision_id"] is not None and int(document_row["active_search_text_revision_id"]) == int(text_revision_id) ) return payload def list_text_revision_summaries_for_document(connection: sqlite3.Connection, document_id: int) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM text_revisions WHERE document_id = ? ORDER BY id DESC """, (document_id,), ).fetchall() return [text_revision_summary_by_id(connection, int(row["id"])) for row in rows] def default_quality_score_for_text_status(text_status: object, text_content: str) -> float | None: normalized_status = normalize_whitespace(str(text_status or "")).lower() if not text_content.strip() or normalized_status == "empty": return 0.0 if normalized_status == "ok": return 1.0 return None TEXT_QUALITY_POOR_THRESHOLD = 0.55 TEXT_QUALITY_IMPROVEMENT_TOLERANCE = 0.02 TRANSLATION_HEADER_LABELS = { "attachments", "bcc", "cc", "date", "from", "sent", "subject", "to", } TRANSLATION_HEADER_PATTERN = re.compile( r"^\s*(Attachments?|Bcc|Cc|Date|From|Sent|Subject|To):\s*(.+?)\s*$", re.IGNORECASE | re.MULTILINE, ) TRANSLATION_EMAIL_PATTERN = re.compile(r"\b[A-Z0-9._%+\-]+@[A-Z0-9.\-]+\.[A-Z]{2,}\b", re.IGNORECASE) TRANSLATION_URL_PATTERN = re.compile(r"\b(?:https?://|www\.)\S+\b", re.IGNORECASE) TRANSLATION_MONEY_PATTERN = re.compile(r"\$\d[\d,]*(?:\.\d+)?") TRANSLATION_BATES_PATTERN = re.compile(r"\b[A-Z]{1,6}\d{5,}\b") TRANSLATION_NUMBER_PATTERN = re.compile(r"\b\d[\d,]*(?:\.\d+)?\b") TRANSLATION_NUMERIC_DATE_PATTERN = re.compile(r"\b\d{1,2}/\d{1,2}/\d{2,4}\b") TRANSLATION_TEXTUAL_DATE_PATTERN = re.compile( r"\b(?:Monday|Tuesday|Wednesday|Thursday|Friday|Saturday|Sunday),?\s+" r"(?:January|February|March|April|May|June|July|August|September|October|November|December)\s+" r"\d{1,2},\s+\d{4}(?:\s+\d{1,2}:\d{2}(?::\d{2})?\s*(?:AM|PM))?\b", re.IGNORECASE, ) TRANSLATION_PHONE_PATTERN = re.compile(r"\b(?:\(\d{3}\)\s*|\d{3}[-.])\d{3}[-.]\d{4}\b") TRANSLATION_FILE_NAME_PATTERN = re.compile(r"\b[^\s;,:()<>]+\.(?:pdf|docx?|xlsx?|pptx?|eml|txt)\b", re.IGNORECASE) COMMON_TEXT_PUNCTUATION = { ".", ",", ";", ":", "!", "?", "(", ")", "[", "]", "{", "}", "<", ">", "/", "\\", "@", "#", "%", "^", "&", "*", "-", "_", "=", "+", "|", "~", "$", "'", '"', "`", "“", "”", "‘", "’", "…", "–", "—", "•", } def ordered_unique_matches( pattern: re.Pattern[str], text: str, *, casefold: bool = False, ) -> list[str]: seen: set[str] = set() matches: list[str] = [] for match in pattern.finditer(text): value = match.group(0) key = value.casefold() if casefold else value if key in seen: continue seen.add(key) matches.append(value) return matches def normalize_translation_header_label(label: str) -> str: normalized = normalize_whitespace(label).lower().rstrip(":") if normalized == "attachment": return "attachments" return normalized def translation_header_rows(text: str) -> list[tuple[str, str]]: return [ (normalize_translation_header_label(match.group(1)), normalize_whitespace(match.group(2))) for match in TRANSLATION_HEADER_PATTERN.finditer(text) if normalize_whitespace(match.group(2)) ] def translation_header_display_names(text: str) -> list[str]: names: list[str] = [] seen: set[str] = set() for label, value in translation_header_rows(text): if label not in {"from", "to", "cc", "bcc"}: continue candidate_values: list[str] = [] parsed_addresses = getaddresses([value]) candidate_values.extend(normalize_whitespace(name) for name, _ in parsed_addresses if normalize_whitespace(name)) for segment in re.split(r";", value): cleaned = normalize_whitespace(re.sub(r"<[^>]+>", "", segment)) if not cleaned or "@" in cleaned or not any(ch.isalpha() for ch in cleaned): continue candidate_values.append(cleaned) for candidate in candidate_values: key = candidate.casefold() if key in seen: continue seen.add(key) names.append(candidate) return names def missing_literal_values(values: list[str], text: str, *, casefold: bool = False) -> list[str]: haystack = text.casefold() if casefold else text missing: list[str] = [] for value in values: needle = value.casefold() if casefold else value if needle not in haystack: missing.append(value) return missing def summarize_missing_literals( source_text: str, translated_text: str, *, label: str, pattern: re.Pattern[str], issues: list[str], casefold: bool = False, ) -> int: source_values = ordered_unique_matches(pattern, source_text, casefold=casefold) missing_values = missing_literal_values(source_values, translated_text, casefold=casefold) if missing_values: sample = ", ".join(missing_values[:3]) suffix = f" (+{len(missing_values) - 3} more)" if len(missing_values) > 3 else "" issues.append(f"{label} changed or disappeared: {sample}{suffix}") return len(source_values) def translation_validation_summary( source_text: str, translated_text: str, *, target_language: str | None, ) -> dict[str, object]: blocking_issues: list[str] = [] warnings: list[str] = [] source_headers = translation_header_rows(source_text) translated_headers = translation_header_rows(translated_text) source_header_labels = [label for label, _ in source_headers if label in TRANSLATION_HEADER_LABELS] translated_header_labels = [label for label, _ in translated_headers if label in TRANSLATION_HEADER_LABELS] if source_header_labels and source_header_labels != translated_header_labels: blocking_issues.append( "Email header labels/order changed: " f"source={source_header_labels}, translated={translated_header_labels}" ) for label, value in source_headers: if label in {"date", "sent"} and value and value not in translated_text: blocking_issues.append(f"Header value for {label!r} changed: {value}") checked_literal_counts = { "emails": summarize_missing_literals( source_text, translated_text, label="Email addresses", pattern=TRANSLATION_EMAIL_PATTERN, issues=blocking_issues, casefold=True, ), "urls": summarize_missing_literals( source_text, translated_text, label="URLs", pattern=TRANSLATION_URL_PATTERN, issues=blocking_issues, casefold=True, ), "money": summarize_missing_literals( source_text, translated_text, label="Currency amounts", pattern=TRANSLATION_MONEY_PATTERN, issues=blocking_issues, ), "bates": summarize_missing_literals( source_text, translated_text, label="Bates numbers", pattern=TRANSLATION_BATES_PATTERN, issues=blocking_issues, ), "numbers": summarize_missing_literals( source_text, translated_text, label="Numbers", pattern=TRANSLATION_NUMBER_PATTERN, issues=blocking_issues, ), "numeric_dates": summarize_missing_literals( source_text, translated_text, label="Numeric dates", pattern=TRANSLATION_NUMERIC_DATE_PATTERN, issues=blocking_issues, ), "textual_dates": summarize_missing_literals( source_text, translated_text, label="Textual dates", pattern=TRANSLATION_TEXTUAL_DATE_PATTERN, issues=blocking_issues, casefold=True, ), "phone_numbers": summarize_missing_literals( source_text, translated_text, label="Phone numbers", pattern=TRANSLATION_PHONE_PATTERN, issues=blocking_issues, ), "file_names": summarize_missing_literals( source_text, translated_text, label="File names", pattern=TRANSLATION_FILE_NAME_PATTERN, issues=blocking_issues, casefold=True, ), } source_names = translation_header_display_names(source_text) missing_names = missing_literal_values(source_names, translated_text, casefold=False) if missing_names: sample = ", ".join(missing_names[:3]) suffix = f" (+{len(missing_names) - 3} more)" if len(missing_names) > 3 else "" blocking_issues.append(f"Header names changed or disappeared: {sample}{suffix}") normalized_source = normalize_whitespace(source_text) normalized_translated = normalize_whitespace(translated_text) if normalized_source and normalized_source == normalized_translated and target_language: warnings.append( f"Translated text is identical to the source despite target_language={target_language!r}." ) score = max(0.0, min(1.0, 1.0 - (0.18 * len(blocking_issues)) - (0.04 * len(warnings)))) return { "kind": "translation_literal_preservation", "target_language": target_language, "status": "failed" if blocking_issues else "ok", "blocking_issues": blocking_issues, "warnings": warnings, "checked_literal_counts": checked_literal_counts, "source_header_labels": source_header_labels, "translated_header_labels": translated_header_labels, "quality_score": score, } def is_common_text_character(value: str) -> bool: if value.isspace() or value.isalnum(): return True category = unicodedata.category(value) if category.startswith(("L", "N")): return True return value in COMMON_TEXT_PUNCTUATION def is_suspicious_text_token(token: str) -> bool: stripped = token.strip() if len(stripped) < 4: return False if "\ufffd" in stripped: return True alpha_count = sum(1 for char in stripped if char.isalpha()) digit_count = sum(1 for char in stripped if char.isdigit()) symbol_count = sum(1 for char in stripped if not is_common_text_character(char)) punctuation_count = sum(1 for char in stripped if not char.isalnum() and not char.isspace()) if symbol_count: return True if alpha_count == 0 and digit_count == 0: return True if alpha_count <= 1 and punctuation_count >= 2: return True if punctuation_count * 2 >= len(stripped): return True return False def text_quality_summary( text_content: str, *, revision_kind: str, source_text: str | None = None, validation_summary: dict[str, object] | None = None, ) -> dict[str, object]: normalized = normalize_whitespace(text_content) stripped = normalized.strip() if not stripped: return { "quality_score": 0.0, "blocking_issues": ["text is empty"], "warnings": [], "metrics": {"char_count": 0, "token_count": 0}, } if revision_kind == "translation" and isinstance(validation_summary, dict): return { "quality_score": float(validation_summary.get("quality_score") or 0.0), "blocking_issues": list(validation_summary.get("blocking_issues") or []), "warnings": list(validation_summary.get("warnings") or []), "metrics": { "char_count": len(text_content), "token_count": len(re.findall(r"\S+", stripped)), }, } tokens = re.findall(r"\S+", stripped) suspicious_tokens = [token for token in tokens if is_suspicious_text_token(token)] suspicious_ratio = len(suspicious_tokens) / max(len(tokens), 1) uncommon_char_count = sum(1 for char in stripped if not is_common_text_character(char)) uncommon_char_ratio = uncommon_char_count / max(len(stripped), 1) replacement_count = stripped.count("\ufffd") blocking_issues: list[str] = [] warnings: list[str] = [] score = 1.0 if revision_kind == "image_description": if len(stripped) < 40: blocking_issues.append("image description is too short to be search-friendly") score -= 0.35 elif len(tokens) < 8: warnings.append("image description is very short") score -= 0.1 else: score -= min(0.45, uncommon_char_ratio * 4.0) score -= min(0.35, suspicious_ratio * 1.8) if replacement_count: blocking_issues.append("text contains replacement characters") score -= 0.2 if source_text: normalized_source = normalize_whitespace(source_text).strip() if normalized_source: source_ratio = len(stripped) / max(len(normalized_source), 1) if len(normalized_source) > 200 and source_ratio < 0.35: warnings.append("text is substantially shorter than the source revision") score -= 0.25 elif len(normalized_source) > 200 and source_ratio < 0.6: warnings.append("text is noticeably shorter than the source revision") score -= 0.1 score = max(0.0, min(1.0, score)) return { "quality_score": score, "blocking_issues": blocking_issues, "warnings": warnings, "metrics": { "char_count": len(text_content), "token_count": len(tokens), "suspicious_token_count": len(suspicious_tokens), "uncommon_char_count": uncommon_char_count, "replacement_count": replacement_count, }, } def text_revision_quality_summary( connection: sqlite3.Connection, paths: dict[str, Path], text_revision_row: sqlite3.Row, ) -> dict[str, object]: text_content = read_text_revision_body(paths, text_revision_row["storage_rel_path"]) or "" provider_metadata = decode_json_text(text_revision_row["provider_metadata_json"], default={}) or {} validation_summary = provider_metadata.get("validation") if isinstance(provider_metadata, dict) else None source_text = None if ( str(text_revision_row["revision_kind"] or "") in {"ocr", "translation"} and text_revision_row["parent_revision_id"] is not None ): parent_row = require_text_revision_row_by_id(connection, int(text_revision_row["parent_revision_id"])) source_text = read_text_revision_body(paths, parent_row["storage_rel_path"]) summary = text_quality_summary( text_content, revision_kind=str(text_revision_row["revision_kind"] or ""), source_text=source_text, validation_summary=validation_summary if isinstance(validation_summary, dict) else None, ) stored_quality_score = ( float(text_revision_row["quality_score"]) if text_revision_row["quality_score"] is not None else None ) effective_quality_score = summary["quality_score"] if stored_quality_score is not None: if str(text_revision_row["revision_kind"] or "") == "source_extract": effective_quality_score = min(stored_quality_score, float(summary["quality_score"])) else: effective_quality_score = stored_quality_score return { "revision_id": int(text_revision_row["id"]), "revision_kind": str(text_revision_row["revision_kind"] or ""), "quality_score": float(summary["quality_score"]), "effective_quality_score": float(effective_quality_score), "stored_quality_score": stored_quality_score, "blocking_issues": list(summary["blocking_issues"]), "warnings": list(summary["warnings"]), "metrics": dict(summary["metrics"]), "text_content": text_content, } def text_revision_activation_comparison( connection: sqlite3.Connection, paths: dict[str, Path], *, document_row: sqlite3.Row, candidate_revision_row: sqlite3.Row, ) -> dict[str, object]: candidate_summary = text_revision_quality_summary(connection, paths, candidate_revision_row) current_summary = None current_revision_row = None if document_row["active_search_text_revision_id"] is not None: current_revision_row = require_text_revision_row_by_id( connection, int(document_row["active_search_text_revision_id"]), ) current_summary = text_revision_quality_summary(connection, paths, current_revision_row) current_char_count = ( len(str(current_summary["text_content"])) if current_summary is not None else 0 ) candidate_char_count = len(str(candidate_summary["text_content"])) return { "candidate_revision_id": int(candidate_revision_row["id"]), "candidate_revision_kind": str(candidate_revision_row["revision_kind"] or ""), "candidate_char_count": candidate_char_count, "candidate_quality_score": candidate_summary["effective_quality_score"], "candidate_blocking_issues": candidate_summary["blocking_issues"], "candidate_warnings": candidate_summary["warnings"], "current_revision_id": ( int(current_revision_row["id"]) if current_revision_row is not None else None ), "current_revision_kind": ( str(current_revision_row["revision_kind"] or "") if current_revision_row is not None else None ), "current_char_count": current_char_count, "current_quality_score": ( current_summary["effective_quality_score"] if current_summary is not None else None ), "current_blocking_issues": ( current_summary["blocking_issues"] if current_summary is not None else [] ), "current_warnings": ( current_summary["warnings"] if current_summary is not None else [] ), "candidate_to_current_char_ratio": ( candidate_char_count / current_char_count if current_char_count > 0 else None ), } def activation_policy_decision( comparison: dict[str, object], *, activation_policy: str, ) -> tuple[bool, str | None]: normalized_policy = normalize_text_revision_activation_policy(activation_policy) if normalized_policy in {"always", "manual"}: return True, None candidate_char_count = int(comparison.get("candidate_char_count") or 0) candidate_quality_score = float(comparison.get("candidate_quality_score") or 0.0) current_char_count = int(comparison.get("current_char_count") or 0) current_quality = comparison.get("current_quality_score") current_quality_score = float(current_quality) if current_quality is not None else None if normalized_policy == "if_empty": if candidate_char_count < 1: return False, "candidate_text_empty" if current_char_count < 1: return True, None return False, "active_text_not_empty" if candidate_char_count < 1: return False, "candidate_text_empty" if current_char_count < 1: return True, None if current_quality_score is None: return True, None if current_quality_score >= TEXT_QUALITY_POOR_THRESHOLD: return False, "active_text_not_poor" if candidate_quality_score + TEXT_QUALITY_IMPROVEMENT_TOLERANCE < current_quality_score: return False, "candidate_quality_regression" return True, None def text_revision_storage_rel_path(document_id: int, revision_kind: str, content_hash: str) -> str: revision_slug = sanitize_processing_identifier(revision_kind, label="Revision kind", prefix="revision") file_name = f"{revision_slug}-{content_hash}.txt" return (Path("text-revisions") / f"doc-{int(document_id):08d}" / file_name).as_posix() def write_text_revision_body(paths: dict[str, Path], storage_rel_path: str, text_content: str) -> None: absolute_path = paths["state_dir"] / storage_rel_path absolute_path.parent.mkdir(parents=True, exist_ok=True) absolute_path.write_text(text_content, encoding="utf-8") def read_text_revision_body(paths: dict[str, Path], storage_rel_path: str | None) -> str | None: if not storage_rel_path: return None absolute_path = paths["state_dir"] / storage_rel_path if not absolute_path.exists(): return None return absolute_path.read_text(encoding="utf-8") def document_row_has_seeded_text_revisions(document_row: sqlite3.Row | None) -> bool: if document_row is None: return False keys = document_row.keys() return ( "source_text_revision_id" in keys and "active_search_text_revision_id" in keys and document_row["source_text_revision_id"] is not None and document_row["active_search_text_revision_id"] is not None ) DOCUMENT_EXTRACTED_METADATA_FIELDS = ( "page_count", "author", "content_type", "date_created", "date_modified", "participants", "title", "subject", "recipients", "text_status", ) def comparable_document_metadata_value(value: object) -> object: if value is None: return None if isinstance(value, (int, float)): return value normalized = normalize_whitespace(str(value)) return normalized or None def extracted_payload_matches_document_row(document_row: sqlite3.Row | None, extracted: object) -> bool: if document_row is None or not isinstance(extracted, dict): return True keys = set(document_row.keys()) if "content_hash" in keys: extracted_content_hash = sha256_text(str(extracted.get("text_content") or "")) if comparable_document_metadata_value(document_row["content_hash"]) != extracted_content_hash: return False for field_name in DOCUMENT_EXTRACTED_METADATA_FIELDS: if field_name not in keys: continue extracted_value = extracted.get("text_status", "ok") if field_name == "text_status" else extracted.get(field_name) if comparable_document_metadata_value(document_row[field_name]) != comparable_document_metadata_value(extracted_value): return False return True def container_documents_missing_text_revisions( connection: sqlite3.Connection, *, source_kind: str, source_rel_path: str, ) -> bool: document_ids = sorted( container_document_ids_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ) if not document_ids: return False placeholders = ", ".join("?" for _ in document_ids) row = connection.execute( f""" SELECT 1 FROM documents WHERE id IN ({placeholders}) AND lifecycle_status != 'deleted' AND ( source_text_revision_id IS NULL OR active_search_text_revision_id IS NULL ) LIMIT 1 """, document_ids, ).fetchone() return row is not None def find_matching_text_revision( connection: sqlite3.Connection, *, document_id: int, revision_kind: str, content_hash: str, parent_revision_id: int | None, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM text_revisions WHERE document_id = ? AND revision_kind = ? AND content_hash = ? AND COALESCE(parent_revision_id, 0) = COALESCE(?, 0) AND retracted_at IS NULL ORDER BY id DESC LIMIT 1 """, (document_id, revision_kind, content_hash, parent_revision_id), ).fetchone() def create_text_revision_row( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, revision_kind: str, text_content: str, language: str | None, parent_revision_id: int | None, created_by_job_version_id: int | None, quality_score: float | None, provider_metadata: dict[str, object] | None, created_at: str | None = None, ) -> int: content_hash = sha256_text(text_content) matching_row = find_matching_text_revision( connection, document_id=document_id, revision_kind=revision_kind, content_hash=content_hash, parent_revision_id=parent_revision_id, ) storage_rel_path = text_revision_storage_rel_path(document_id, revision_kind, content_hash) write_text_revision_body(paths, storage_rel_path, text_content) if matching_row is not None: return int(matching_row["id"]) timestamp = created_at or utc_now() cursor = connection.execute( """ INSERT INTO text_revisions ( document_id, revision_kind, language, parent_revision_id, created_by_job_version_id, storage_rel_path, content_hash, char_count, token_estimate, quality_score, provider_metadata_json, created_at, retracted_at, retraction_reason ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL) """, ( document_id, revision_kind, language, parent_revision_id, created_by_job_version_id, storage_rel_path, content_hash, len(text_content), token_estimate(text_content) if text_content.strip() else 0, quality_score, compact_json_text(provider_metadata or {}), timestamp, ), ) return int(cursor.lastrowid) def ensure_source_text_revision_for_document( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, text_content: str, text_status: object, existing_source_revision_id: int | None = None, language: str | None = None, created_at: str | None = None, ) -> int: content_hash = sha256_text(text_content) if existing_source_revision_id is not None: existing_row = connection.execute( """ SELECT * FROM text_revisions WHERE id = ? AND document_id = ? AND revision_kind = 'source_extract' AND retracted_at IS NULL """, (existing_source_revision_id, document_id), ).fetchone() if existing_row is not None and str(existing_row["content_hash"]) == content_hash: storage_rel_path = str(existing_row["storage_rel_path"] or "") if storage_rel_path: write_text_revision_body(paths, storage_rel_path, text_content) return int(existing_row["id"]) return create_text_revision_row( connection, paths, document_id=document_id, revision_kind="source_extract", text_content=text_content, language=language, parent_revision_id=None, created_by_job_version_id=None, quality_score=default_quality_score_for_text_status(text_status, text_content), provider_metadata={"text_status": text_status}, created_at=created_at, ) def set_document_text_revision_pointers( connection: sqlite3.Connection, *, document_id: int, source_text_revision_id: int, active_search_text_revision_id: int | None = None, ) -> None: source_row = connection.execute( """ SELECT revision_kind, language, quality_score FROM text_revisions WHERE id = ? """, (source_text_revision_id,), ).fetchone() if source_row is None: raise RetrieverError(f"Unknown source text revision id: {source_text_revision_id}") active_revision_id = active_search_text_revision_id or source_text_revision_id active_row = connection.execute( """ SELECT revision_kind, language, quality_score FROM text_revisions WHERE id = ? """, (active_revision_id,), ).fetchone() if active_row is None: raise RetrieverError(f"Unknown active text revision id: {active_revision_id}") connection.execute( """ UPDATE documents SET source_text_revision_id = ?, active_search_text_revision_id = ?, active_text_source_kind = ?, active_text_language = ?, active_text_quality_score = ? WHERE id = ? """, ( source_text_revision_id, active_revision_id, active_row["revision_kind"], active_row["language"], active_row["quality_score"], document_id, ), ) def seed_source_text_revision_for_document( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, extracted: dict[str, object], existing_row: sqlite3.Row | None = None, created_at: str | None = None, ) -> int: text_content = str(extracted.get("text_content") or "") revision_id = ensure_source_text_revision_for_document( connection, paths, document_id=document_id, text_content=text_content, text_status=extracted.get("text_status"), existing_source_revision_id=( int(existing_row["source_text_revision_id"]) if existing_row is not None and "source_text_revision_id" in existing_row.keys() and existing_row["source_text_revision_id"] is not None else None ), language=(str(extracted.get("language")) if extracted.get("language") else None), created_at=created_at, ) set_document_text_revision_pointers( connection, document_id=document_id, source_text_revision_id=revision_id, active_search_text_revision_id=revision_id, ) return revision_id def record_text_revision_activation_event( connection: sqlite3.Connection, *, document_id: int, text_revision_id: int, activated_by_job_version_id: int | None, source_result_id: int | None, activation_policy: str, created_at: str | None = None, ) -> int: timestamp = created_at or utc_now() cursor = connection.execute( """ INSERT INTO text_revision_activation_events ( document_id, text_revision_id, activated_by_job_version_id, source_result_id, activation_policy, created_at ) VALUES (?, ?, ?, ?, ?, ?) """, ( document_id, text_revision_id, activated_by_job_version_id, source_result_id, activation_policy, timestamp, ), ) return int(cursor.lastrowid) def activate_text_revision_for_document( connection: sqlite3.Connection, paths: dict[str, Path], *, document_id: int, text_revision_id: int, activation_policy: str = "manual", activated_by_job_version_id: int | None = None, source_result_id: int | None = None, ) -> dict[str, object]: normalized_policy = normalize_text_revision_activation_policy(activation_policy) document_row = connection.execute( """ SELECT id, source_text_revision_id, active_search_text_revision_id, content_type, conversation_id, production_id FROM documents WHERE id = ? """, (document_id,), ).fetchone() if document_row is None: raise RetrieverError(f"Unknown document id: {document_id}") text_revision_row = require_text_revision_row_by_id(connection, text_revision_id) if int(text_revision_row["document_id"]) != int(document_id): raise RetrieverError( f"Text revision {text_revision_id} belongs to document {text_revision_row['document_id']}, " f"not document {document_id}." ) text_content = read_text_revision_body(paths, text_revision_row["storage_rel_path"]) if text_content is None: raise RetrieverError(f"Text revision {text_revision_id} has no readable body on disk.") comparison = text_revision_activation_comparison( connection, paths, document_row=document_row, candidate_revision_row=text_revision_row, ) should_activate, skip_reason = activation_policy_decision( comparison, activation_policy=normalized_policy, ) if not should_activate: return { "status": "skipped", "document_id": int(document_id), "text_revision": text_revision_summary_by_id(connection, int(text_revision_id)), "activation_policy": normalized_policy, "skip_reason": skip_reason, "comparison": comparison, "preview_regen": None, } replace_document_chunks(connection, document_id, chunk_text(text_content)) source_text_revision_id = ( int(document_row["source_text_revision_id"]) if document_row["source_text_revision_id"] is not None else root_text_revision_id(connection, text_revision_id) ) set_document_text_revision_pointers( connection, document_id=document_id, source_text_revision_id=source_text_revision_id, active_search_text_revision_id=int(text_revision_id), ) now = utc_now() normalized_text = normalize_whitespace(text_content) connection.execute( """ UPDATE documents SET content_hash = ?, text_status = ?, updated_at = ? WHERE id = ? """, ( text_revision_row["content_hash"], "ok" if normalized_text else "empty", now, document_id, ), ) activation_event_id = record_text_revision_activation_event( connection, document_id=document_id, text_revision_id=int(text_revision_id), activated_by_job_version_id=activated_by_job_version_id, source_result_id=source_result_id, activation_policy=normalized_policy, created_at=now, ) # Best-effort: regenerate the preview so the on-disk HTML reflects the # newly active text. Failures never roll back activation. try: if document_content_type_is_chat(document_row["content_type"]): if document_row["conversation_id"] is not None: preview_regen = { "status": "ok", "conversation_id": int(document_row["conversation_id"]), "refreshed_conversations": refresh_conversation_previews( connection, paths, [int(document_row["conversation_id"])], ), } else: preview_regen = regenerate_chat_preview_for_document( connection, paths, document_id=document_id, ) elif document_content_type_is_email(document_row["content_type"]): if document_row["conversation_id"] is not None: preview_regen = { "status": "ok", "conversation_id": int(document_row["conversation_id"]), "refreshed_conversations": refresh_conversation_previews( connection, paths, [int(document_row["conversation_id"])], ), } elif document_row["production_id"] is None: preview_regen = regenerate_email_preview_for_document( connection, paths, document_id=document_id, ) else: preview_regen = regenerate_production_preview_for_document( connection, paths, document_id=document_id, text_content=text_content, ) else: preview_regen = regenerate_production_preview_for_document( connection, paths, document_id=document_id, text_content=text_content, ) except Exception as exc: # noqa: BLE001 - best-effort regen preview_regen = {"status": "failed", "error": str(exc)} return { "status": "ok", "document_id": int(document_id), "text_revision": text_revision_summary_by_id(connection, int(text_revision_id)), "activation_event_id": activation_event_id, "activation_policy": normalized_policy, "comparison": comparison, "preview_regen": preview_regen, } def run_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "job_version_id": int(row["job_version_id"]), "from_run_id": row["from_run_id"], "selector": decode_json_text(row["selector_json"], default={}) or {}, "exclude_selector": decode_json_text(row["exclude_selector_json"], default={}) or {}, "activation_policy": str(row["activation_policy"] or "manual"), "family_mode": row["family_mode"], "seed_limit": row["seed_limit"], "status": row["status"], "planned_count": int(row["planned_count"] or 0), "completed_count": int(row["completed_count"] or 0), "failed_count": int(row["failed_count"] or 0), "skipped_count": int(row["skipped_count"] or 0), "created_at": row["created_at"], "started_at": row["started_at"], "completed_at": row["completed_at"], "canceled_at": row["canceled_at"], } def run_snapshot_document_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_id": int(row["run_id"]), "document_id": int(row["document_id"]), "ordinal": int(row["ordinal"]), "inclusion_reason": decode_json_text(row["inclusion_reason_json"], default={}) or {}, "pinned_input_revision_id": row["pinned_input_revision_id"], "pinned_input_identity": row["pinned_input_identity"], "pinned_content_hash": row["pinned_content_hash"], "created_at": row["created_at"], } def create_run_row( connection: sqlite3.Connection, *, job_version_id: int, selector: dict[str, object], exclude_selector: dict[str, object], activation_policy: str, family_mode: str, seed_limit: int | None, from_run_id: int | None, status: str = "planned", ) -> int: now = utc_now() cursor = connection.execute( """ INSERT INTO runs ( job_version_id, from_run_id, selector_json, exclude_selector_json, activation_policy, family_mode, seed_limit, status, planned_count, completed_count, failed_count, skipped_count, created_at, started_at, completed_at, canceled_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0, 0, 0, 0, ?, NULL, NULL, NULL) """, ( job_version_id, from_run_id, compact_json_text(selector), compact_json_text(exclude_selector), activation_policy, family_mode, seed_limit, status, now, ), ) return int(cursor.lastrowid) def maybe_activate_created_text_revision( connection: sqlite3.Connection, paths: dict[str, Path], *, run_row: sqlite3.Row, job_version_row: sqlite3.Row, document_id: int, result_id: int, text_revision_id: int | None, ) -> dict[str, object] | None: if text_revision_id is None: return None activation_policy = normalize_run_activation_policy(str(run_row["activation_policy"] or "manual")) if activation_policy == "manual": return None activation_payload = activate_text_revision_for_document( connection, paths, document_id=document_id, text_revision_id=text_revision_id, activation_policy=activation_policy, activated_by_job_version_id=int(job_version_row["id"]), source_result_id=result_id, ) if str(activation_payload.get("status") or "") != "ok": return None return activation_payload def replace_run_snapshot_documents( connection: sqlite3.Connection, *, run_id: int, snapshot_rows: list[dict[str, object]], ) -> None: connection.execute("DELETE FROM run_snapshot_documents WHERE run_id = ?", (run_id,)) if snapshot_rows: connection.executemany( """ INSERT INTO run_snapshot_documents ( run_id, document_id, ordinal, inclusion_reason_json, pinned_input_revision_id, pinned_input_identity, pinned_content_hash, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, [ ( run_id, row["document_id"], row["ordinal"], compact_json_text(row["inclusion_reason"]), row["pinned_input_revision_id"], row["pinned_input_identity"], row["pinned_content_hash"], row.get("created_at", utc_now()), ) for row in snapshot_rows ], ) connection.execute( """ UPDATE runs SET planned_count = ? WHERE id = ? """, (len(snapshot_rows), run_id), ) def run_item_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_id": int(row["run_id"]), "run_snapshot_document_id": row["run_snapshot_document_id"], "item_kind": row["item_kind"], "document_id": int(row["document_id"]), "page_number": row["page_number"], "segment_id": row["segment_id"], "input_artifact_rel_path": row["input_artifact_rel_path"], "input_identity": row["input_identity"], "result_id": row["result_id"], "status": row["status"], "claimed_by": row["claimed_by"], "claimed_at": row["claimed_at"], "last_heartbeat_at": row["last_heartbeat_at"], "attempt_count": int(row["attempt_count"] or 0), "last_error": row["last_error"], "created_at": row["created_at"], "started_at": row["started_at"], "completed_at": row["completed_at"], } def find_run_item_row( connection: sqlite3.Connection, *, run_id: int, run_snapshot_document_id: int | None, item_kind: str, page_number: int | None, segment_id: int | None, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM run_items WHERE run_id = ? AND ((run_snapshot_document_id = ?) OR (run_snapshot_document_id IS NULL AND ? IS NULL)) AND item_kind = ? AND ((page_number = ?) OR (page_number IS NULL AND ? IS NULL)) AND ((segment_id = ?) OR (segment_id IS NULL AND ? IS NULL)) ORDER BY id ASC LIMIT 1 """, ( run_id, run_snapshot_document_id, run_snapshot_document_id, item_kind, page_number, page_number, segment_id, segment_id, ), ).fetchone() def ensure_run_item_row( connection: sqlite3.Connection, *, run_id: int, run_snapshot_document_id: int | None, item_kind: str, document_id: int, input_identity: str, page_number: int | None = None, segment_id: int | None = None, input_artifact_rel_path: str | None = None, ) -> sqlite3.Row: normalized_item_kind = normalize_run_item_kind(item_kind) existing_row = find_run_item_row( connection, run_id=run_id, run_snapshot_document_id=run_snapshot_document_id, item_kind=normalized_item_kind, page_number=page_number, segment_id=segment_id, ) if existing_row is not None: if ( str(existing_row["input_identity"]) != input_identity or normalize_whitespace(str(existing_row["input_artifact_rel_path"] or "")) != normalize_whitespace(str(input_artifact_rel_path or "")) ): connection.execute( """ UPDATE run_items SET input_identity = ?, input_artifact_rel_path = ? WHERE id = ? """, (input_identity, input_artifact_rel_path, existing_row["id"]), ) return connection.execute("SELECT * FROM run_items WHERE id = ?", (existing_row["id"],)).fetchone() return existing_row now = utc_now() connection.execute( """ INSERT OR IGNORE INTO run_items ( run_id, run_snapshot_document_id, item_kind, document_id, page_number, segment_id, input_artifact_rel_path, input_identity, result_id, status, claimed_by, claimed_at, last_heartbeat_at, attempt_count, last_error, created_at, started_at, completed_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, NULL, 'pending', NULL, NULL, NULL, 0, NULL, ?, NULL, NULL) """, ( run_id, run_snapshot_document_id, normalized_item_kind, document_id, page_number, segment_id, input_artifact_rel_path, input_identity, now, ), ) ensured_row = find_run_item_row( connection, run_id=run_id, run_snapshot_document_id=run_snapshot_document_id, item_kind=normalized_item_kind, page_number=page_number, segment_id=segment_id, ) assert ensured_row is not None return ensured_row def update_run_item_row( connection: sqlite3.Connection, *, run_item_id: int, status: str, result_id: int | None = None, last_error: str | None = None, claimed_by: str | None = None, claimed_at: str | None = None, last_heartbeat_at: str | None = None, started_at: str | None = None, completed_at: str | None = None, increment_attempt_count: bool = False, ) -> None: normalized_status = normalize_run_item_status(status) row = connection.execute( """ SELECT attempt_count, claimed_by, claimed_at, last_heartbeat_at, started_at, completed_at FROM run_items WHERE id = ? """, (run_item_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown run item id: {run_item_id}") next_attempt_count = int(row["attempt_count"] or 0) + (1 if increment_attempt_count else 0) effective_claimed_by = claimed_by if claimed_by is not None else row["claimed_by"] effective_claimed_at = claimed_at if claimed_at is not None else row["claimed_at"] effective_last_heartbeat_at = last_heartbeat_at if last_heartbeat_at is not None else row["last_heartbeat_at"] effective_started_at = started_at if started_at is not None else row["started_at"] effective_completed_at = completed_at if completed_at is not None else row["completed_at"] connection.execute( """ UPDATE run_items SET status = ?, result_id = ?, claimed_by = ?, claimed_at = ?, last_heartbeat_at = ?, attempt_count = ?, last_error = ?, started_at = ?, completed_at = ? WHERE id = ? """, ( normalized_status, result_id, effective_claimed_by, effective_claimed_at, effective_last_heartbeat_at, next_attempt_count, last_error, effective_started_at, effective_completed_at, run_item_id, ), ) def list_run_item_payloads_for_run(connection: sqlite3.Connection, run_id: int) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM run_items WHERE run_id = ? ORDER BY id ASC """, (run_id,), ).fetchall() return [run_item_row_to_payload(row) for row in rows] def run_worker_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_id": int(row["run_id"]), "claimed_by": row["claimed_by"], "launch_mode": row["launch_mode"], "worker_task_id": row["worker_task_id"], "status": row["status"], "max_batches": row["max_batches"], "batches_prepared": int(row["batches_prepared"] or 0), "items_completed": int(row["items_completed"] or 0), "items_failed": int(row["items_failed"] or 0), "last_heartbeat_at": row["last_heartbeat_at"], "last_error": row["last_error"], "created_at": row["created_at"], "started_at": row["started_at"], "completed_at": row["completed_at"], "cancel_requested_at": row["cancel_requested_at"], "summary": decode_json_text(row["summary_json"], default={}) or {}, } def find_run_worker_row( connection: sqlite3.Connection, *, run_id: int, claimed_by: str, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM run_workers WHERE run_id = ? AND claimed_by = ? """, (run_id, claimed_by), ).fetchone() def ensure_run_worker_row( connection: sqlite3.Connection, *, run_id: int, claimed_by: str, launch_mode: str, worker_task_id: str | None, max_batches: int | None, ) -> sqlite3.Row: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") normalized_launch_mode = normalize_run_worker_mode(launch_mode) normalized_worker_task_id = ( normalize_whitespace(worker_task_id) if worker_task_id is not None else "" ) or None existing_row = find_run_worker_row(connection, run_id=run_id, claimed_by=normalized_claimed_by) now = utc_now() if existing_row is not None: updates: list[str] = [] params: list[object] = [] if str(existing_row["launch_mode"]) != normalized_launch_mode: updates.append("launch_mode = ?") params.append(normalized_launch_mode) if normalized_worker_task_id and normalize_whitespace(str(existing_row["worker_task_id"] or "")) != normalized_worker_task_id: updates.append("worker_task_id = ?") params.append(normalized_worker_task_id) if max_batches is not None and existing_row["max_batches"] != max_batches: updates.append("max_batches = ?") params.append(max_batches) if str(existing_row["status"] or "") in {"completed", "failed", "stopped", "orphaned", "canceled"}: updates.append("status = 'active'") updates.append("completed_at = NULL") updates.append("last_error = NULL") updates.append("cancel_requested_at = NULL") updates.append("summary_json = '{}'") if updates: connection.execute( f""" UPDATE run_workers SET {", ".join(updates)} WHERE id = ? """, (*params, existing_row["id"]), ) return connection.execute("SELECT * FROM run_workers WHERE id = ?", (existing_row["id"],)).fetchone() connection.execute( """ INSERT INTO run_workers ( run_id, claimed_by, launch_mode, worker_task_id, status, max_batches, batches_prepared, items_completed, items_failed, last_heartbeat_at, last_error, created_at, started_at, completed_at, cancel_requested_at, summary_json ) VALUES (?, ?, ?, ?, 'active', ?, 0, 0, 0, NULL, NULL, ?, ?, NULL, NULL, '{}') """, ( run_id, normalized_claimed_by, normalized_launch_mode, normalized_worker_task_id, max_batches, now, now, ), ) return find_run_worker_row(connection, run_id=run_id, claimed_by=normalized_claimed_by) def update_run_worker_row( connection: sqlite3.Connection, *, run_id: int, claimed_by: str, status: str | None = None, worker_task_id: str | None = None, max_batches: int | None = None, increment_batches_prepared: bool = False, increment_items_completed: int = 0, increment_items_failed: int = 0, heartbeat: bool = False, last_error: str | None = None, cancel_requested: bool = False, summary: dict[str, object] | None = None, completed_at: str | None = None, ) -> sqlite3.Row | None: row = find_run_worker_row(connection, run_id=run_id, claimed_by=claimed_by) if row is None: return None next_status = normalize_run_worker_status(status) if status is not None else str(row["status"]) next_batches_prepared = int(row["batches_prepared"] or 0) + (1 if increment_batches_prepared else 0) next_items_completed = int(row["items_completed"] or 0) + max(0, int(increment_items_completed)) next_items_failed = int(row["items_failed"] or 0) + max(0, int(increment_items_failed)) next_worker_task_id = ( normalize_whitespace(worker_task_id) if worker_task_id is not None else normalize_whitespace(str(row["worker_task_id"] or "")) ) or row["worker_task_id"] next_max_batches = max_batches if max_batches is not None else row["max_batches"] next_last_heartbeat_at = utc_now() if heartbeat else row["last_heartbeat_at"] next_cancel_requested_at = utc_now() if cancel_requested else row["cancel_requested_at"] next_summary = summary if summary is not None else (decode_json_text(row["summary_json"], default={}) or {}) next_completed_at = completed_at if completed_at is not None else row["completed_at"] connection.execute( """ UPDATE run_workers SET status = ?, worker_task_id = ?, max_batches = ?, batches_prepared = ?, items_completed = ?, items_failed = ?, last_heartbeat_at = ?, last_error = ?, completed_at = ?, cancel_requested_at = ?, summary_json = ? WHERE id = ? """, ( next_status, next_worker_task_id, next_max_batches, next_batches_prepared, next_items_completed, next_items_failed, next_last_heartbeat_at, last_error if last_error is not None else row["last_error"], next_completed_at, next_cancel_requested_at, compact_json_text(next_summary), row["id"], ), ) return connection.execute("SELECT * FROM run_workers WHERE id = ?", (row["id"],)).fetchone() def list_run_worker_payloads_for_run(connection: sqlite3.Connection, run_id: int) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM run_workers WHERE run_id = ? ORDER BY created_at ASC, id ASC """, (run_id,), ).fetchall() return [run_worker_row_to_payload(row) for row in rows] def require_run_row_by_id(connection: sqlite3.Connection, run_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM runs WHERE id = ? """, (run_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown run id: {run_id}") return row def require_run_item_row_by_id(connection: sqlite3.Connection, run_item_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM run_items WHERE id = ? """, (run_item_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown run item id: {run_item_id}") return row def workspace_relative_artifact_path(paths: dict[str, Path], absolute_path: Path) -> str: try: relative_to_state = absolute_path.relative_to(paths["state_dir"]) except ValueError: return relative_document_path(paths["root"], absolute_path) return str(Path(INTERNAL_REL_PATH_PREFIX) / relative_to_state) def resolve_workspace_artifact_path(root: Path, rel_path: str | None) -> Path | None: normalized = normalize_whitespace(str(rel_path or "")) if not normalized: return None path = Path(normalized) if path.parts and path.parts[0] == INTERNAL_REL_PATH_PREFIX: # Synthetic rel_paths mirror the state-directory layout. state_relative = Path(*path.parts[1:]) if len(path.parts) > 1 else Path() return (root / ".retriever" / state_relative).resolve() return (root / normalized).resolve() def freeze_ocr_source_artifact( paths: dict[str, Path], root: Path, *, source_rel_path: str, run_id: int, document_id: int, page_number: int, ) -> str: source_path = resolve_workspace_artifact_path(root, source_rel_path) if source_path is None or not source_path.exists(): raise RetrieverError(f"OCR source artifact is missing: {source_rel_path!r}") output_dir = paths["jobs_dir"] / "ocr" / f"run-{run_id}" / f"doc-{document_id}" output_dir.mkdir(parents=True, exist_ok=True) output_name = f"page-{page_number:04d}-{source_path.name}" output_path = output_dir / output_name if not output_path.exists(): output_path.write_bytes(source_path.read_bytes()) return workspace_relative_artifact_path(paths, output_path) def render_pdf_pages_for_ocr( paths: dict[str, Path], *, source_path: Path, run_id: int, document_id: int, resolution: int, ) -> list[dict[str, object]]: pdfplumber_module = dependency_guard("pdfplumber", "pdfplumber", "pdf") output_dir = paths["jobs_dir"] / "ocr" / f"run-{run_id}" / f"doc-{document_id}" output_dir.mkdir(parents=True, exist_ok=True) page_specs: list[dict[str, object]] = [] with pdfplumber_module.open(source_path) as pdf: # type: ignore[union-attr] for page_number, page in enumerate(pdf.pages, start=1): output_path = output_dir / f"page-{page_number:04d}.png" if not output_path.exists(): page_image = page.to_image(resolution=resolution) page_image.save(output_path, format="PNG") page_specs.append( { "page_number": page_number, "input_artifact_rel_path": workspace_relative_artifact_path(paths, output_path), } ) return page_specs def ocr_page_item_specs_for_document( connection: sqlite3.Connection, paths: dict[str, Path], root: Path, *, run_id: int, document_row: sqlite3.Row, job_version_row: sqlite3.Row, ) -> list[dict[str, object]]: input_basis = str(job_version_row["input_basis"]) if input_basis == "source_parts": source_part_rows = connection.execute( """ SELECT part_kind, rel_source_path, ordinal FROM document_source_parts WHERE document_id = ? AND part_kind IN ('image', 'native') ORDER BY CASE part_kind WHEN 'image' THEN 0 WHEN 'native' THEN 1 ELSE 2 END ASC, ordinal ASC, id ASC """, (document_row["id"],), ).fetchall() image_part_rows = [row for row in source_part_rows if str(row["part_kind"]) == "image"] if image_part_rows: return [ { "page_number": int(row["ordinal"] or index), "input_artifact_rel_path": freeze_ocr_source_artifact( paths, root, source_rel_path=str(row["rel_source_path"]), run_id=run_id, document_id=int(document_row["id"]), page_number=int(row["ordinal"] or index), ), } for index, row in enumerate(image_part_rows, start=1) ] native_part_row = next((row for row in source_part_rows if str(row["part_kind"]) == "native"), None) if native_part_row is not None: native_source_rel_path = str(native_part_row["rel_source_path"]) native_source_path = resolve_workspace_artifact_path(root, native_source_rel_path) if native_source_path is None or not native_source_path.exists(): raise RetrieverError( f"Document {document_row['id']} has no accessible native source part for OCR: {native_source_rel_path!r}." ) native_file_type = normalize_extension(native_source_path) if native_file_type == "pdf": parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} rendering_settings = parameters.get("rendering_settings") if isinstance(parameters, dict) else None resolution = DEFAULT_OCR_RENDER_RESOLUTION if isinstance(rendering_settings, dict) and rendering_settings.get("resolution") is not None: try: resolution = max(72, int(rendering_settings["resolution"])) except (TypeError, ValueError): resolution = DEFAULT_OCR_RENDER_RESOLUTION return render_pdf_pages_for_ocr( paths, source_path=native_source_path, run_id=run_id, document_id=int(document_row["id"]), resolution=resolution, ) if native_file_type in IMAGE_NATIVE_PREVIEW_FILE_TYPES: return [ { "page_number": 1, "input_artifact_rel_path": freeze_ocr_source_artifact( paths, root, source_rel_path=native_source_rel_path, run_id=run_id, document_id=int(document_row["id"]), page_number=1, ), } ] raise RetrieverError( f"Document {document_row['id']} native source part is not OCR-capable: {native_file_type!r}." ) raise RetrieverError( f"Document {document_row['id']} has source parts but no image/native OCR source parts." ) source_path = resolve_workspace_artifact_path(root, str(document_row["rel_path"])) if source_path is None or not source_path.exists(): raise RetrieverError(f"Document {document_row['id']} has no accessible source path for OCR.") file_type = normalize_extension(source_path) or normalize_whitespace(str(document_row["file_type"] or "")).lower() if file_type == "pdf": parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} rendering_settings = parameters.get("rendering_settings") if isinstance(parameters, dict) else None resolution = DEFAULT_OCR_RENDER_RESOLUTION if isinstance(rendering_settings, dict) and rendering_settings.get("resolution") is not None: try: resolution = max(72, int(rendering_settings["resolution"])) except (TypeError, ValueError): resolution = DEFAULT_OCR_RENDER_RESOLUTION return render_pdf_pages_for_ocr( paths, source_path=source_path, run_id=run_id, document_id=int(document_row["id"]), resolution=resolution, ) if file_type in IMAGE_NATIVE_PREVIEW_FILE_TYPES: return [ { "page_number": 1, "input_artifact_rel_path": freeze_ocr_source_artifact( paths, root, source_rel_path=str(document_row["rel_path"]), run_id=run_id, document_id=int(document_row["id"]), page_number=1, ), } ] raise RetrieverError( f"OCR page materialization does not know how to prepare document {document_row['id']} with file type {file_type!r}." ) def prior_run_page_item_specs( connection: sqlite3.Connection, *, from_run_id: int, document_id: int, ) -> list[dict[str, object]]: rows = connection.execute( """ SELECT page_number, input_artifact_rel_path FROM run_items WHERE run_id = ? AND document_id = ? AND item_kind = 'page' AND input_artifact_rel_path IS NOT NULL ORDER BY page_number ASC, id ASC """, (from_run_id, document_id), ).fetchall() return [ { "page_number": int(row["page_number"] or 0), "input_artifact_rel_path": str(row["input_artifact_rel_path"]), } for row in rows ] def materialize_run_items_for_run( connection: sqlite3.Connection, paths: dict[str, Path], root: Path, run_id: int, ) -> list[sqlite3.Row]: existing_rows = connection.execute( """ SELECT * FROM run_items WHERE run_id = ? ORDER BY id ASC """, (run_id,), ).fetchall() if existing_rows: return existing_rows run_row = require_run_row_by_id(connection, run_id) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( """ SELECT * FROM jobs WHERE id = ? """, (job_version_row["job_id"],), ).fetchone() assert job_row is not None job_kind = normalize_job_kind(str(job_row["job_kind"])) snapshot_rows = connection.execute( """ SELECT * FROM run_snapshot_documents WHERE run_id = ? ORDER BY ordinal ASC, id ASC """, (run_id,), ).fetchall() materialized_rows: list[sqlite3.Row] = [] for snapshot_row in snapshot_rows: document_row = connection.execute( """ SELECT * FROM documents WHERE id = ? """, (snapshot_row["document_id"],), ).fetchone() if document_row is None: continue if job_kind in {"ocr", "image_description"}: from_run_id = int(run_row["from_run_id"]) if run_row["from_run_id"] is not None else None page_specs = ( prior_run_page_item_specs( connection, from_run_id=from_run_id, document_id=int(snapshot_row["document_id"]), ) if from_run_id is not None else [] ) if not page_specs: page_specs = ocr_page_item_specs_for_document( connection, paths, root, run_id=run_id, document_row=document_row, job_version_row=job_version_row, ) for page_spec in page_specs: materialized_rows.append( ensure_run_item_row( connection, run_id=run_id, run_snapshot_document_id=int(snapshot_row["id"]), item_kind="page", document_id=int(snapshot_row["document_id"]), page_number=int(page_spec["page_number"]), input_artifact_rel_path=str(page_spec["input_artifact_rel_path"]), input_identity=str(snapshot_row["pinned_input_identity"]), ) ) continue materialized_rows.append( ensure_run_item_row( connection, run_id=run_id, run_snapshot_document_id=int(snapshot_row["id"]), item_kind="document", document_id=int(snapshot_row["document_id"]), input_identity=str(snapshot_row["pinned_input_identity"]), ) ) return materialized_rows def reuse_active_results_for_run(connection: sqlite3.Connection, run_id: int) -> int: run_row = require_run_row_by_id(connection, run_id) job_version_id = int(run_row["job_version_id"]) pending_rows = connection.execute( """ SELECT * FROM run_items WHERE run_id = ? AND status = 'pending' ORDER BY id ASC """, (run_id,), ).fetchall() reused_count = 0 for row in pending_rows: existing_result_row = find_active_result_row( connection, document_id=int(row["document_id"]), job_version_id=job_version_id, input_identity=str(row["input_identity"]), ) if existing_result_row is None: continue update_run_item_row( connection, run_item_id=int(row["id"]), status="skipped", result_id=int(existing_result_row["id"]), last_error=None, completed_at=utc_now(), ) reused_count += 1 return reused_count def ocr_page_output_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_item_id": int(row["run_item_id"]), "run_id": int(row["run_id"]), "document_id": int(row["document_id"]), "page_number": int(row["page_number"]), "text_content": row["text_content"], "raw_output": decode_json_text(row["raw_output_json"]), "normalized_output": decode_json_text(row["normalized_output_json"]), "provider_metadata": decode_json_text(row["provider_metadata_json"], default={}) or {}, "created_at": row["created_at"], } def find_ocr_page_output_row(connection: sqlite3.Connection, *, run_item_id: int) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM ocr_page_outputs WHERE run_item_id = ? """, (run_item_id,), ).fetchone() def upsert_ocr_page_output_row( connection: sqlite3.Connection, *, run_item_id: int, run_id: int, document_id: int, page_number: int, text_content: str, raw_output: object, normalized_output: object, provider_metadata: dict[str, object] | None, ) -> tuple[int, bool]: existing_row = find_ocr_page_output_row(connection, run_item_id=run_item_id) if existing_row is not None: return int(existing_row["id"]), False cursor = connection.execute( """ INSERT INTO ocr_page_outputs ( run_item_id, run_id, document_id, page_number, text_content, raw_output_json, normalized_output_json, provider_metadata_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_item_id, run_id, document_id, page_number, text_content, compact_json_text(raw_output), compact_json_text(normalized_output), compact_json_text(provider_metadata or {}), utc_now(), ), ) return int(cursor.lastrowid), True def list_ocr_page_output_payloads_for_document( connection: sqlite3.Connection, *, run_id: int, document_id: int, ) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM ocr_page_outputs WHERE run_id = ? AND document_id = ? ORDER BY page_number ASC, id ASC """, (run_id, document_id), ).fetchall() return [ocr_page_output_row_to_payload(row) for row in rows] def image_description_page_output_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_item_id": int(row["run_item_id"]), "run_id": int(row["run_id"]), "document_id": int(row["document_id"]), "page_number": int(row["page_number"]), "text_content": row["text_content"], "raw_output": decode_json_text(row["raw_output_json"]), "normalized_output": decode_json_text(row["normalized_output_json"]), "provider_metadata": decode_json_text(row["provider_metadata_json"], default={}) or {}, "created_at": row["created_at"], } def find_image_description_page_output_row( connection: sqlite3.Connection, *, run_item_id: int, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM image_description_page_outputs WHERE run_item_id = ? """, (run_item_id,), ).fetchone() def upsert_image_description_page_output_row( connection: sqlite3.Connection, *, run_item_id: int, run_id: int, document_id: int, page_number: int, text_content: str, raw_output: object, normalized_output: object, provider_metadata: dict[str, object] | None, ) -> tuple[int, bool]: existing_row = find_image_description_page_output_row(connection, run_item_id=run_item_id) if existing_row is not None: return int(existing_row["id"]), False cursor = connection.execute( """ INSERT INTO image_description_page_outputs ( run_item_id, run_id, document_id, page_number, text_content, raw_output_json, normalized_output_json, provider_metadata_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_item_id, run_id, document_id, page_number, text_content, compact_json_text(raw_output), compact_json_text(normalized_output), compact_json_text(provider_metadata or {}), utc_now(), ), ) return int(cursor.lastrowid), True def list_image_description_page_output_payloads_for_document( connection: sqlite3.Connection, *, run_id: int, document_id: int, ) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM image_description_page_outputs WHERE run_id = ? AND document_id = ? ORDER BY page_number ASC, id ASC """, (run_id, document_id), ).fetchall() return [image_description_page_output_row_to_payload(row) for row in rows] def claimed_run_item_rows_for_worker( connection: sqlite3.Connection, *, run_id: int, claimed_by: str, ) -> list[sqlite3.Row]: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") return connection.execute( """ SELECT * FROM run_items WHERE run_id = ? AND status = 'running' AND result_id IS NULL AND claimed_by = ? ORDER BY COALESCE(page_number, 0) ASC, id ASC """, (run_id, normalized_claimed_by), ).fetchall() def claim_run_item_rows( connection: sqlite3.Connection, *, run_id: int, claimed_by: str, limit: int, stale_after_seconds: int = DEFAULT_RUN_ITEM_CLAIM_STALE_SECONDS, ) -> list[sqlite3.Row]: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") if limit < 1: return [] run_row = require_run_row_by_id(connection, run_id) if str(run_row["status"] or "") in {"canceled", "completed", "failed"}: return [] stale_cutoff = format_utc_timestamp(datetime.now(timezone.utc) - timedelta(seconds=max(1, stale_after_seconds))) stale_claimed_rows = connection.execute( """ SELECT DISTINCT claimed_by FROM run_items WHERE run_id = ? AND status = 'running' AND result_id IS NULL AND claimed_by IS NOT NULL AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ? """, (run_id, stale_cutoff), ).fetchall() for stale_row in stale_claimed_rows: stale_claimed_by = normalize_whitespace(str(stale_row["claimed_by"] or "")) if not stale_claimed_by or stale_claimed_by == normalized_claimed_by: continue update_run_worker_row( connection, run_id=run_id, claimed_by=stale_claimed_by, status="orphaned", last_error="Worker claim heartbeat expired and items were reclaimed.", completed_at=utc_now(), ) candidate_rows = connection.execute( """ SELECT * FROM run_items WHERE run_id = ? AND ( status = 'pending' OR ( status = 'running' AND result_id IS NULL AND claimed_by IS NOT NULL AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ? ) ) ORDER BY id ASC LIMIT ? """, (run_id, stale_cutoff, limit), ).fetchall() if not candidate_rows: return [] now = utc_now() item_ids = [int(row["id"]) for row in candidate_rows] placeholders = ", ".join("?" for _ in item_ids) connection.execute( f""" UPDATE run_items SET status = 'running', claimed_by = ?, claimed_at = ?, last_heartbeat_at = ?, started_at = COALESCE(started_at, ?), completed_at = NULL WHERE id IN ({placeholders}) """, (normalized_claimed_by, now, now, now, *item_ids), ) connection.execute( """ UPDATE runs SET status = 'running', started_at = COALESCE(started_at, ?), completed_at = NULL, canceled_at = NULL WHERE id = ? """, (now, run_id), ) return connection.execute( f""" SELECT * FROM run_items WHERE id IN ({placeholders}) ORDER BY id ASC """, item_ids, ).fetchall() def heartbeat_claimed_run_items( connection: sqlite3.Connection, *, run_id: int, claimed_by: str, ) -> int: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") now = utc_now() cursor = connection.execute( """ UPDATE run_items SET last_heartbeat_at = ? WHERE run_id = ? AND status = 'running' AND claimed_by = ? """, (now, run_id, normalized_claimed_by), ) update_run_worker_row( connection, run_id=run_id, claimed_by=normalized_claimed_by, heartbeat=True, ) return int(cursor.rowcount or 0) def cancel_pending_run_items(connection: sqlite3.Connection, *, run_id: int) -> int: now = utc_now() cursor = connection.execute( """ UPDATE run_items SET status = 'skipped', last_error = COALESCE(last_error, 'Run canceled.'), completed_at = COALESCE(completed_at, ?) WHERE run_id = ? AND status = 'pending' """, (now, run_id), ) return int(cursor.rowcount or 0) def request_run_worker_cancellation( connection: sqlite3.Connection, *, run_id: int, force: bool, ) -> list[str]: worker_rows = connection.execute( """ SELECT claimed_by, worker_task_id FROM run_workers WHERE run_id = ? AND status = 'active' ORDER BY claimed_by ASC """, (run_id,), ).fetchall() task_ids: list[str] = [] for row in worker_rows: claimed_by = normalize_whitespace(str(row["claimed_by"] or "")) worker_task_id = normalize_whitespace(str(row["worker_task_id"] or "")) or None if worker_task_id: task_ids.append(worker_task_id) update_run_worker_row( connection, run_id=run_id, claimed_by=claimed_by, status="canceled", cancel_requested=True, last_error=("Force-stop requested." if force else "Cancellation requested."), completed_at=utc_now() if force else None, ) return task_ids def run_item_status_counts(connection: sqlite3.Connection, run_id: int) -> dict[str, int]: return { str(row["status"]): int(row["count"] or 0) for row in connection.execute( """ SELECT status, COUNT(*) AS count FROM run_items WHERE run_id = ? GROUP BY status """, (run_id,), ).fetchall() } def run_execution_metadata(connection: sqlite3.Connection, run_id: int) -> tuple[sqlite3.Row, sqlite3.Row, sqlite3.Row]: run_row = require_run_row_by_id(connection, run_id) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( """ SELECT * FROM jobs WHERE id = ? """, (job_version_row["job_id"],), ).fetchone() if job_row is None: raise RetrieverError(f"Run {run_id} references a missing job id: {job_version_row['job_id']}") return run_row, job_version_row, job_row def ocr_run_pending_finalization_count(connection: sqlite3.Connection, run_id: int) -> int: row = connection.execute( """ SELECT COUNT(*) AS count FROM run_snapshot_documents AS snapshot WHERE snapshot.run_id = ? AND EXISTS ( SELECT 1 FROM run_items AS page_item WHERE page_item.run_id = snapshot.run_id AND page_item.document_id = snapshot.document_id AND page_item.item_kind = 'page' ) AND NOT EXISTS ( SELECT 1 FROM run_items AS finalized_item WHERE finalized_item.run_id = snapshot.run_id AND finalized_item.document_id = snapshot.document_id AND finalized_item.result_id IS NOT NULL ) """, (run_id,), ).fetchone() return int(row["count"] or 0) def image_description_run_pending_finalization_count(connection: sqlite3.Connection, run_id: int) -> int: row = connection.execute( """ SELECT COUNT(*) AS count FROM run_snapshot_documents AS snapshot WHERE snapshot.run_id = ? AND EXISTS ( SELECT 1 FROM run_items AS page_item WHERE page_item.run_id = snapshot.run_id AND page_item.document_id = snapshot.document_id AND page_item.item_kind = 'page' ) AND NOT EXISTS ( SELECT 1 FROM run_items AS finalized_item WHERE finalized_item.run_id = snapshot.run_id AND finalized_item.document_id = snapshot.document_id AND finalized_item.result_id IS NOT NULL ) """, (run_id,), ).fetchone() return int(row["count"] or 0) RUN_WORKER_HANDOFF_SUFFIX_PATTERN = re.compile(r"^(?P.+?)-handoff(?:-(?P\d+))?$") def handoff_claimed_by_base(run_id: int, claimed_by: str | None) -> str: normalized = normalize_whitespace(claimed_by or "") or f"cowork-run-{int(run_id)}" match = RUN_WORKER_HANDOFF_SUFFIX_PATTERN.fullmatch(normalized) if match is not None: return normalize_whitespace(match.group("base") or "") or normalized return normalized def next_handoff_claimed_by_hint( connection: sqlite3.Connection, *, run_id: int, claimed_by: str | None, ) -> str: base = handoff_claimed_by_base(run_id, claimed_by) existing_claimed_by = { normalize_whitespace(str(row["claimed_by"] or "")) for row in connection.execute( """ SELECT claimed_by FROM run_workers WHERE run_id = ? AND claimed_by IS NOT NULL UNION SELECT claimed_by FROM run_items WHERE run_id = ? AND claimed_by IS NOT NULL """, (run_id, run_id), ).fetchall() if normalize_whitespace(str(row["claimed_by"] or "")) } first_candidate = f"{base}-handoff" if first_candidate not in existing_claimed_by: return first_candidate for index in range(2, 1000): candidate = f"{base}-handoff-{index}" if candidate not in existing_claimed_by: return candidate return f"{base}-handoff-{secrets.token_hex(3)}" def build_run_worker_payload( connection: sqlite3.Connection, run_id: int, *, run_payload: dict[str, object], claimed_by: str | None = None, ) -> dict[str, object]: _, job_version_row, job_row = run_execution_metadata(connection, run_id) job_kind = normalize_job_kind(str(job_row["job_kind"])) capability = normalize_job_capability(str(job_version_row["capability"])) run_item_counts = dict(run_payload.get("run_item_counts") or {}) pending_count = int(run_item_counts.get("pending", 0) or 0) running_count = int(run_item_counts.get("running", 0) or 0) completed_count = int(run_item_counts.get("completed", 0) or 0) failed_count = int(run_item_counts.get("failed", 0) or 0) skipped_count = int(run_item_counts.get("skipped", 0) or 0) planned_count = int(run_payload.get("planned_count", 0) or 0) total_items = pending_count + running_count + completed_count + failed_count + skipped_count outstanding_items = max(planned_count - completed_count - failed_count - skipped_count, 0) needs_ocr_finalization = job_kind == "ocr" and ocr_run_pending_finalization_count(connection, run_id) > 0 needs_image_description_finalization = ( job_kind == "image_description" and image_description_run_pending_finalization_count(connection, run_id) > 0 ) normalized_claimed_by = normalize_whitespace(claimed_by) if claimed_by and claimed_by.strip() else None worker_row = ( find_run_worker_row(connection, run_id=run_id, claimed_by=normalized_claimed_by) if normalized_claimed_by else None ) claimed_rows_for_worker = ( claimed_run_item_rows_for_worker(connection, run_id=run_id, claimed_by=normalized_claimed_by) if normalized_claimed_by else [] ) claimed_item_count = len(claimed_rows_for_worker) recommended_execution_mode = ( "background" if max(total_items, planned_count) > DEFAULT_WORKER_INLINE_MAX_ITEMS else "inline" ) recommended_batch_size = min( DEFAULT_RUN_ITEM_CLAIM_BATCH_SIZE, DEFAULT_WORKER_BATCH_SIZE, max(outstanding_items, 1), ) recommended_max_batches = ( DEFAULT_WORKER_BACKGROUND_MAX_BATCHES if recommended_execution_mode == "background" else DEFAULT_WORKER_INLINE_MAX_BATCHES ) effective_max_batches = ( int(worker_row["max_batches"]) if worker_row is not None and worker_row["max_batches"] is not None else recommended_max_batches ) worker_batches_prepared = int(worker_row["batches_prepared"] or 0) if worker_row is not None else 0 worker_cancel_requested = worker_row is not None and worker_row["cancel_requested_at"] is not None worker_should_handoff = ( worker_row is not None and effective_max_batches > 0 and worker_batches_prepared >= effective_max_batches and pending_count > 0 ) handoff_claimed_by_hint = ( next_handoff_claimed_by_hint(connection, run_id=run_id, claimed_by=normalized_claimed_by) if worker_should_handoff else None ) effective_launch_mode = ( normalize_run_worker_mode(str(worker_row["launch_mode"])) if worker_row is not None and worker_row["launch_mode"] is not None else recommended_execution_mode ) claim_stale_after_seconds = default_run_item_claim_stale_seconds_for_launch_mode(effective_launch_mode) recommended_heartbeat_interval_seconds = max(10, claim_stale_after_seconds // 2) next_action = "claim" stop_reason = None run_status = str(run_payload.get("status") or "") if run_status == "canceled": next_action = "stop" stop_reason = "canceled" elif worker_cancel_requested: next_action = "stop" stop_reason = "canceled" elif claimed_item_count > 0: next_action = "process_batch" elif worker_should_handoff: next_action = "handoff" stop_reason = "max_batches_reached" elif needs_ocr_finalization and pending_count == 0 and running_count == 0 and failed_count == 0: next_action = "finalize_ocr" elif needs_image_description_finalization and pending_count == 0 and running_count == 0 and failed_count == 0: next_action = "finalize_image_description" elif planned_count == 0: next_action = "stop" stop_reason = "empty" elif outstanding_items == 0 and pending_count == 0 and running_count == 0: next_action = "stop" if failed_count: stop_reason = "failed" elif run_status == "completed": stop_reason = "completed" else: stop_reason = run_status or "idle" return { "job_kind": job_kind, "capability": capability, "claimed_by": normalized_claimed_by, "launch_mode": worker_row["launch_mode"] if worker_row is not None else None, "worker_task_id": worker_row["worker_task_id"] if worker_row is not None else None, "worker_status": worker_row["status"] if worker_row is not None else None, "recommended_execution_mode": recommended_execution_mode, "recommended_batch_size": recommended_batch_size, "recommended_max_batches_per_worker": recommended_max_batches, "max_batches_per_worker": effective_max_batches, "claimed_item_count": claimed_item_count, "claim_stale_after_seconds": claim_stale_after_seconds, "recommended_heartbeat_interval_seconds": recommended_heartbeat_interval_seconds, "batches_prepared": worker_batches_prepared, "outstanding_items": outstanding_items, "needs_ocr_finalization": needs_ocr_finalization, "needs_image_description_finalization": needs_image_description_finalization, "handoff_required": worker_should_handoff, "handoff_claimed_by_hint": handoff_claimed_by_hint, "handoff_message": ( "This worker reached max_batches_per_worker; finish it and continue with a new --claimed-by identity." if worker_should_handoff else None ), "should_exit_after_batch": worker_should_handoff, "after_batch_action": "handoff" if worker_should_handoff else None, "next_action": next_action, "stop_reason": stop_reason, } def recent_run_item_failures( connection: sqlite3.Connection, *, run_id: int, limit: int = 10, ) -> list[dict[str, object]]: rows = connection.execute( """ SELECT id, document_id, segment_id, last_error, completed_at FROM run_items WHERE run_id = ? AND status = 'failed' ORDER BY completed_at DESC, id DESC LIMIT ? """, (run_id, max(1, limit)), ).fetchall() return [ { "run_item_id": int(row["id"]), "document_id": int(row["document_id"]), "segment_id": row["segment_id"], "error": row["last_error"], "completed_at": row["completed_at"], } for row in rows ] def build_run_supervision_payload( connection: sqlite3.Connection, run_id: int, *, run_payload: dict[str, object], worker_payloads: list[dict[str, object]], ) -> dict[str, object]: run_status = str(run_payload.get("status") or "") run_worker_hint = build_run_worker_payload(connection, run_id, run_payload=run_payload) active_workers = [payload for payload in worker_payloads if str(payload["status"]) == "active"] background_workers = [payload for payload in active_workers if str(payload["launch_mode"]) == "background"] orphaned_workers = [payload for payload in worker_payloads if str(payload["status"]) == "orphaned"] max_parallel_workers = DEFAULT_WORKER_BACKGROUND_MAX_PARALLEL outstanding_items = int(run_worker_hint["outstanding_items"] or 0) finalization_pending = bool( run_worker_hint["needs_ocr_finalization"] or run_worker_hint["needs_image_description_finalization"] ) if run_worker_hint["recommended_execution_mode"] == "background" and outstanding_items > 0: suggested_worker_count = min( max_parallel_workers, max(1, (outstanding_items + DEFAULT_WORKER_BATCH_SIZE - 1) // DEFAULT_WORKER_BATCH_SIZE), ) elif outstanding_items > 0: suggested_worker_count = 1 else: suggested_worker_count = 0 spawn_additional_worker_count = max(0, suggested_worker_count - len(active_workers)) force_stop_task_ids = [ str(payload["worker_task_id"]) for payload in worker_payloads if payload["cancel_requested_at"] is not None and payload["worker_task_id"] ] should_schedule_wakeup = bool( run_status not in {"canceled", "completed", "failed"} and (outstanding_items > 0 or len(active_workers) > 0 or finalization_pending) ) if run_status in {"canceled", "completed", "failed"}: wakeup_reason = None elif finalization_pending and not active_workers: wakeup_reason = "finalization_pending" elif len(active_workers) > 0: wakeup_reason = "workers_active" elif outstanding_items > 0: wakeup_reason = "pending_work" else: wakeup_reason = None recommended_action = "wait" if run_status == "canceled": recommended_action = "stop" elif run_worker_hint["next_action"] in {"finalize_ocr", "finalize_image_description"}: recommended_action = str(run_worker_hint["next_action"]) elif run_worker_hint["next_action"] == "stop": recommended_action = "stop" elif spawn_additional_worker_count > 0: recommended_action = "spawn_background_worker" if run_worker_hint["recommended_execution_mode"] == "background" else "claim_inline" return { "active_worker_count": len(active_workers), "background_worker_count": len(background_workers), "orphaned_worker_count": len(orphaned_workers), "continuation_needed": bool(outstanding_items and not active_workers and run_status not in {"canceled", "completed", "failed"}), "outstanding_items": outstanding_items, "finalization_pending": finalization_pending, "max_parallel_workers": max_parallel_workers, "suggested_worker_count": suggested_worker_count, "spawn_additional_worker_count": spawn_additional_worker_count, "should_schedule_wakeup": should_schedule_wakeup, "wake_interval_seconds": (DEFAULT_WORKER_BACKGROUND_WAKE_INTERVAL_SECONDS if should_schedule_wakeup else None), "wakeup_reason": wakeup_reason, "force_stop_task_ids": force_stop_task_ids, "recommended_action": recommended_action, } def run_item_claim_health( connection: sqlite3.Connection, run_id: int, *, stale_after_seconds: int = DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS, ) -> dict[str, object]: now = datetime.now(timezone.utc) stale_cutoff = format_utc_timestamp(now - timedelta(seconds=max(1, int(stale_after_seconds)))) def age_seconds(raw_value: object) -> int | None: parsed = parse_utc_timestamp(raw_value) if parsed is None: return None return max(0, int((now - parsed).total_seconds())) active_row = connection.execute( """ SELECT MIN(last_heartbeat_at) AS oldest_heartbeat_at, COUNT(*) AS count FROM run_items WHERE run_id = ? AND status = 'running' AND claimed_by IS NOT NULL AND last_heartbeat_at IS NOT NULL """, (run_id,), ).fetchone() stale_row = connection.execute( """ SELECT MIN(last_heartbeat_at) AS oldest_heartbeat_at, COUNT(*) AS count FROM run_items WHERE run_id = ? AND status = 'running' AND claimed_by IS NOT NULL AND last_heartbeat_at IS NOT NULL AND last_heartbeat_at < ? """, (run_id, stale_cutoff), ).fetchone() return { "stale_after_seconds": int(stale_after_seconds), "active_claim_count": int(active_row["count"] or 0) if active_row is not None else 0, "stale_claim_count": int(stale_row["count"] or 0) if stale_row is not None else 0, "oldest_active_claim_heartbeat_at": active_row["oldest_heartbeat_at"] if active_row is not None else None, "oldest_active_claim_age_seconds": ( age_seconds(active_row["oldest_heartbeat_at"]) if active_row is not None else None ), "oldest_stale_claim_heartbeat_at": stale_row["oldest_heartbeat_at"] if stale_row is not None else None, "oldest_stale_claim_age_seconds": ( age_seconds(stale_row["oldest_heartbeat_at"]) if stale_row is not None else None ), "recoverable_by_another_call": bool(int(stale_row["count"] or 0) if stale_row is not None else 0), } def attempt_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_item_id": int(row["run_item_id"]), "attempt_number": int(row["attempt_number"]), "provider_request_id": row["provider_request_id"], "input_tokens": row["input_tokens"], "output_tokens": row["output_tokens"], "cost_cents": row["cost_cents"], "latency_ms": row["latency_ms"], "provider_metadata": decode_json_text(row["provider_metadata_json"], default={}) or {}, "error_summary": row["error_summary"], "created_at": row["created_at"], } def next_attempt_number_for_run_item(connection: sqlite3.Connection, run_item_id: int) -> int: row = connection.execute( """ SELECT COALESCE(MAX(attempt_number), 0) AS max_attempt_number FROM attempts WHERE run_item_id = ? """, (run_item_id,), ).fetchone() return int(row["max_attempt_number"] or 0) + 1 def create_attempt_row( connection: sqlite3.Connection, *, run_item_id: int, provider_request_id: str | None, input_tokens: int | None, output_tokens: int | None, cost_cents: int | None, latency_ms: int | None, provider_metadata: dict[str, object] | None, error_summary: str | None, ) -> int: attempt_number = next_attempt_number_for_run_item(connection, run_item_id) cursor = connection.execute( """ INSERT INTO attempts ( run_item_id, attempt_number, provider_request_id, input_tokens, output_tokens, cost_cents, latency_ms, provider_metadata_json, error_summary, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_item_id, attempt_number, provider_request_id, input_tokens, output_tokens, cost_cents, latency_ms, compact_json_text(provider_metadata or {}), error_summary, utc_now(), ), ) return int(cursor.lastrowid) def result_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "run_id": row["run_id"], "document_id": int(row["document_id"]), "job_version_id": int(row["job_version_id"]), "input_revision_id": row["input_revision_id"], "input_identity": row["input_identity"], "raw_output": decode_json_text(row["raw_output_json"]), "normalized_output": decode_json_text(row["normalized_output_json"]), "created_text_revision_id": row["created_text_revision_id"], "provider_metadata": decode_json_text(row["provider_metadata_json"], default={}) or {}, "created_at": row["created_at"], "retracted_at": row["retracted_at"], "retraction_reason": row["retraction_reason"], } def find_active_result_row( connection: sqlite3.Connection, *, document_id: int, job_version_id: int, input_identity: str, ) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM results WHERE document_id = ? AND job_version_id = ? AND input_identity = ? AND retracted_at IS NULL ORDER BY id DESC LIMIT 1 """, (document_id, job_version_id, input_identity), ).fetchone() def create_result_row( connection: sqlite3.Connection, *, run_id: int, document_id: int, job_version_id: int, input_revision_id: int | None, input_identity: str, raw_output: object, normalized_output: object, created_text_revision_id: int | None, provider_metadata: dict[str, object] | None, ) -> tuple[int, bool]: existing_row = find_active_result_row( connection, document_id=document_id, job_version_id=job_version_id, input_identity=input_identity, ) if existing_row is not None: return int(existing_row["id"]), False now = utc_now() try: cursor = connection.execute( """ INSERT INTO results ( run_id, document_id, job_version_id, input_revision_id, input_identity, raw_output_json, normalized_output_json, created_text_revision_id, provider_metadata_json, created_at, retracted_at, retraction_reason ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL) """, ( run_id, document_id, job_version_id, input_revision_id, input_identity, compact_json_text(raw_output), compact_json_text(normalized_output), created_text_revision_id, compact_json_text(provider_metadata or {}), now, ), ) return int(cursor.lastrowid), True except sqlite3.IntegrityError: existing_row = find_active_result_row( connection, document_id=document_id, job_version_id=job_version_id, input_identity=input_identity, ) if existing_row is None: raise return int(existing_row["id"]), False def result_output_display_value(output_value: object) -> str | None: if output_value is None: return None if isinstance(output_value, bool): return "true" if output_value else "false" if isinstance(output_value, (int, float)): return str(output_value) if isinstance(output_value, str): return output_value return compact_json_text(output_value) def upsert_result_output_rows( connection: sqlite3.Connection, *, result_id: int, job_output_rows: list[sqlite3.Row], output_values_by_name: dict[str, object], ) -> None: now = utc_now() for job_output_row in job_output_rows: output_name = str(job_output_row["output_name"]) output_value = output_values_by_name.get(output_name) connection.execute( """ INSERT INTO result_outputs ( result_id, job_output_id, output_value_json, display_value, score, created_at ) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(result_id, job_output_id) DO NOTHING """, ( result_id, int(job_output_row["id"]), compact_json_text(output_value), result_output_display_value(output_value), None, now, ), ) def result_output_row_to_payload(row: sqlite3.Row) -> dict[str, object]: return { "id": int(row["id"]), "result_id": int(row["result_id"]), "job_output_id": int(row["job_output_id"]), "output_name": row["output_name"], "value_type": row["value_type"], "bound_custom_field": row["bound_custom_field"], "output_value": decode_json_text(row["output_value_json"]), "display_value": row["display_value"], "score": row["score"], "created_at": row["created_at"], } def result_outputs_for_result(connection: sqlite3.Connection, result_id: int) -> list[dict[str, object]]: rows = connection.execute( """ SELECT ro.*, jo.output_name, jo.value_type, jo.bound_custom_field FROM result_outputs ro JOIN job_outputs jo ON jo.id = ro.job_output_id WHERE ro.result_id = ? ORDER BY jo.ordinal ASC, jo.output_name ASC, ro.id ASC """, (result_id,), ).fetchall() return [result_output_row_to_payload(row) for row in rows] def result_summary_by_id(connection: sqlite3.Connection, result_id: int) -> dict[str, object]: row = connection.execute( """ SELECT * FROM results WHERE id = ? """, (result_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown result id: {result_id}") payload = result_row_to_payload(row) payload["outputs"] = result_outputs_for_result(connection, result_id) return payload def list_result_summaries( connection: sqlite3.Connection, *, run_id: int | None = None, document_id: int | None = None, ) -> list[dict[str, object]]: if run_id is not None: rows = connection.execute( """ SELECT DISTINCT r.* FROM run_items ri JOIN results r ON r.id = ri.result_id WHERE ri.run_id = ? ORDER BY r.id DESC """, (run_id,), ).fetchall() if document_id is not None: rows = [row for row in rows if int(row["document_id"]) == int(document_id)] return [result_summary_by_id(connection, int(row["id"])) for row in rows] if document_id is None: rows = connection.execute( """ SELECT * FROM results ORDER BY id DESC """ ).fetchall() else: rows = connection.execute( """ SELECT * FROM results WHERE document_id = ? ORDER BY id DESC """, (document_id,), ).fetchall() return [result_summary_by_id(connection, int(row["id"])) for row in rows] def refresh_run_progress(connection: sqlite3.Connection, run_id: int) -> None: snapshot_count = int( connection.execute( """ SELECT COUNT(*) AS count FROM run_snapshot_documents WHERE run_id = ? """, (run_id,), ).fetchone()["count"] or 0 ) run_item_total = int( connection.execute( """ SELECT COUNT(*) AS count FROM run_items WHERE run_id = ? """, (run_id,), ).fetchone()["count"] or 0 ) planned_count = run_item_total or snapshot_count counts = { str(row["status"]): int(row["count"] or 0) for row in connection.execute( """ SELECT status, COUNT(*) AS count FROM run_items WHERE run_id = ? GROUP BY status """, (run_id,), ).fetchall() } completed_count = counts.get("completed", 0) failed_count = counts.get("failed", 0) skipped_count = counts.get("skipped", 0) running_count = counts.get("running", 0) pending_count = counts.get("pending", 0) row = connection.execute( """ SELECT started_at, canceled_at FROM runs WHERE id = ? """, (run_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown run id: {run_id}") if row["canceled_at"] is not None: status = "canceled" elif planned_count == 0: status = "completed" elif completed_count + failed_count + skipped_count >= planned_count and pending_count == 0 and running_count == 0: status = "failed" if failed_count else "completed" elif running_count > 0: status = "running" else: status = "planned" started_at = row["started_at"] or (utc_now() if status in {"running", "completed", "failed"} else None) completed_at = utc_now() if status in {"completed", "failed", "canceled"} else None connection.execute( """ UPDATE runs SET status = ?, planned_count = ?, completed_count = ?, failed_count = ?, skipped_count = ?, started_at = ?, completed_at = ? WHERE id = ? """, ( status, planned_count, completed_count, failed_count, skipped_count, started_at, completed_at, run_id, ), ) def run_status_by_id( connection: sqlite3.Connection, run_id: int, *, claimed_by: str | None = None, ) -> dict[str, object]: refresh_run_progress(connection, run_id) payload = run_summary_by_id(connection, run_id) status_counts = run_item_status_counts(connection, run_id) payload["run_item_counts"] = { "pending": status_counts.get("pending", 0), "running": status_counts.get("running", 0), "completed": status_counts.get("completed", 0), "failed": status_counts.get("failed", 0), "skipped": status_counts.get("skipped", 0), } payload["snapshot_document_count"] = len(payload.get("documents", [])) payload["recent_failures"] = recent_run_item_failures(connection, run_id=run_id) claim_health = run_item_claim_health(connection, run_id) now = datetime.now(timezone.utc) def claim_age_seconds(raw_value: object) -> int | None: parsed = parse_utc_timestamp(raw_value) if parsed is None: return None return max(0, int((now - parsed).total_seconds())) active_claim_rows = connection.execute( """ SELECT claimed_by, id, document_id, item_kind, page_number, segment_id, last_heartbeat_at FROM run_items WHERE run_id = ? AND status = 'running' AND claimed_by IS NOT NULL ORDER BY claimed_by ASC, id ASC """, (run_id,), ).fetchall() active_claims_by_worker: dict[str, dict[str, object]] = {} for row in active_claim_rows: worker_claimed_by = normalize_whitespace(str(row["claimed_by"] or "")) if not worker_claimed_by: continue claim_payload = active_claims_by_worker.setdefault( worker_claimed_by, { "claimed_by": worker_claimed_by, "count": 0, "run_item_ids": [], "document_ids": [], "page_numbers": [], "segment_ids": [], "last_heartbeat_at": None, }, ) claim_payload["count"] = int(claim_payload["count"] or 0) + 1 claim_payload["run_item_ids"].append(int(row["id"])) document_id = int(row["document_id"]) if document_id not in claim_payload["document_ids"]: claim_payload["document_ids"].append(document_id) if row["page_number"] is not None: claim_payload["page_numbers"].append(int(row["page_number"])) if row["segment_id"] is not None and row["segment_id"] not in claim_payload["segment_ids"]: claim_payload["segment_ids"].append(row["segment_id"]) candidate_heartbeat = row["last_heartbeat_at"] if candidate_heartbeat is not None and ( claim_payload["last_heartbeat_at"] is None or str(candidate_heartbeat) > str(claim_payload["last_heartbeat_at"]) ): claim_payload["last_heartbeat_at"] = candidate_heartbeat payload["active_claims"] = [] for worker_claim in active_claims_by_worker.values(): last_heartbeat_at = worker_claim["last_heartbeat_at"] last_heartbeat_age_seconds = claim_age_seconds(last_heartbeat_at) payload["active_claims"].append( { **worker_claim, "last_heartbeat_at": last_heartbeat_at, "last_heartbeat_age_seconds": last_heartbeat_age_seconds, "stale": ( last_heartbeat_age_seconds is not None and int(last_heartbeat_age_seconds or 0) > int(claim_health["stale_after_seconds"]) ), } ) payload["claim_health"] = claim_health payload["workers"] = list_run_worker_payloads_for_run(connection, run_id) payload["worker"] = build_run_worker_payload( connection, run_id, run_payload=payload, claimed_by=claimed_by, ) payload["supervision"] = build_run_supervision_payload( connection, run_id, run_payload=payload, worker_payloads=payload["workers"], ) return payload def finalize_ocr_results_for_run( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: int, ) -> dict[str, object]: run_row = require_run_row_by_id(connection, run_id) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( """ SELECT * FROM jobs WHERE id = ? """, (job_version_row["job_id"],), ).fetchone() assert job_row is not None if normalize_job_kind(str(job_row["job_kind"])) != "ocr": raise RetrieverError("finalize-ocr-run only supports OCR jobs.") pending_counts = run_item_status_counts(connection, run_id) if pending_counts.get("pending", 0) or pending_counts.get("running", 0) or pending_counts.get("failed", 0): raise RetrieverError( "OCR run cannot be finalized until all page items are completed. " f"pending={pending_counts.get('pending', 0)}, " f"running={pending_counts.get('running', 0)}, " f"failed={pending_counts.get('failed', 0)}." ) snapshot_rows = connection.execute( """ SELECT * FROM run_snapshot_documents WHERE run_id = ? ORDER BY ordinal ASC, id ASC """, (run_id,), ).fetchall() finalized_results: list[dict[str, object]] = [] activations: list[dict[str, object]] = [] for snapshot_row in snapshot_rows: document_id = int(snapshot_row["document_id"]) existing_result_row = find_active_result_row( connection, document_id=document_id, job_version_id=int(job_version_row["id"]), input_identity=str(snapshot_row["pinned_input_identity"]), ) if existing_result_row is not None and existing_result_row["created_text_revision_id"] is not None: result_id = int(existing_result_row["id"]) connection.execute( """ UPDATE run_items SET result_id = ? WHERE run_id = ? AND document_id = ? """, (result_id, run_id, document_id), ) activation_payload = maybe_activate_created_text_revision( connection, paths, run_row=run_row, job_version_row=job_version_row, document_id=document_id, result_id=result_id, text_revision_id=int(existing_result_row["created_text_revision_id"]), ) if activation_payload is not None: activations.append(activation_payload) finalized_results.append(result_summary_by_id(connection, result_id)) continue page_rows = connection.execute( """ SELECT * FROM run_items WHERE run_id = ? AND document_id = ? AND item_kind = 'page' ORDER BY page_number ASC, id ASC """, (run_id, document_id), ).fetchall() if not page_rows: raise RetrieverError(f"OCR run {run_id} has no page items for document {document_id}.") page_output_payloads = list_ocr_page_output_payloads_for_document( connection, run_id=run_id, document_id=document_id, ) if len(page_output_payloads) != len(page_rows): raise RetrieverError( f"OCR run {run_id} is missing page outputs for document {document_id}: " f"expected {len(page_rows)}, found {len(page_output_payloads)}." ) merged_text = "\n\n".join(str(payload["text_content"]) for payload in page_output_payloads if str(payload["text_content"]).strip()) parent_revision_id = ( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row["pinned_input_revision_id"] is not None else None ) parent_text = None if parent_revision_id is not None: parent_revision_row = require_text_revision_row_by_id(connection, parent_revision_id) parent_text = read_text_revision_body(paths, parent_revision_row["storage_rel_path"]) quality_summary = text_quality_summary( merged_text, revision_kind="ocr", source_text=parent_text, ) created_text_revision_id = create_text_revision_row( connection, paths, document_id=document_id, revision_kind="ocr", text_content=merged_text, language=None, parent_revision_id=parent_revision_id, created_by_job_version_id=int(job_version_row["id"]), quality_score=float(quality_summary["quality_score"]), provider_metadata={ "run_id": run_id, "page_count": len(page_output_payloads), "finalized_from": "ocr_page_outputs", "quality": quality_summary, }, ) result_id, _ = create_result_row( connection, run_id=run_id, document_id=document_id, job_version_id=int(job_version_row["id"]), input_revision_id=( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row["pinned_input_revision_id"] is not None else None ), input_identity=str(snapshot_row["pinned_input_identity"]), raw_output={ "page_count": len(page_output_payloads), "finalized_from": "ocr_page_outputs", }, normalized_output={ "page_count": len(page_output_payloads), "created_text_revision_id": created_text_revision_id, }, created_text_revision_id=created_text_revision_id, provider_metadata={ "run_id": run_id, "page_count": len(page_output_payloads), "quality": quality_summary, }, ) connection.execute( """ UPDATE run_items SET result_id = ? WHERE run_id = ? AND document_id = ? """, (result_id, run_id, document_id), ) activation_payload = maybe_activate_created_text_revision( connection, paths, run_row=run_row, job_version_row=job_version_row, document_id=document_id, result_id=result_id, text_revision_id=created_text_revision_id, ) if activation_payload is not None: activations.append(activation_payload) finalized_results.append(result_summary_by_id(connection, result_id)) refresh_run_progress(connection, run_id) stale_text_derived_metadata_document_ids = sorted({int(payload["document_id"]) for payload in activations}) return { "status": "ok", "run": run_status_by_id(connection, run_id), "results": finalized_results, "activations": activations, "stale_text_derived_metadata_document_ids": stale_text_derived_metadata_document_ids, } def build_image_description_revision_text(page_output_payloads: list[dict[str, object]]) -> str: sections: list[str] = [] for payload in page_output_payloads: page_number = int(payload["page_number"]) text_content = normalize_whitespace(str(payload["text_content"] or "")) if not text_content: continue sections.append(f"[IMAGE DESCRIPTION - PAGE {page_number}]\n{text_content}") return "\n\n".join(sections) def finalize_image_description_results_for_run( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: int, ) -> dict[str, object]: run_row = require_run_row_by_id(connection, run_id) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( """ SELECT * FROM jobs WHERE id = ? """, (job_version_row["job_id"],), ).fetchone() assert job_row is not None if normalize_job_kind(str(job_row["job_kind"])) != "image_description": raise RetrieverError("finalize-image-description-run only supports image_description jobs.") pending_counts = run_item_status_counts(connection, run_id) if pending_counts.get("pending", 0) or pending_counts.get("running", 0) or pending_counts.get("failed", 0): raise RetrieverError( "Image-description run cannot be finalized until all page items are completed. " f"pending={pending_counts.get('pending', 0)}, " f"running={pending_counts.get('running', 0)}, " f"failed={pending_counts.get('failed', 0)}." ) snapshot_rows = connection.execute( """ SELECT * FROM run_snapshot_documents WHERE run_id = ? ORDER BY ordinal ASC, id ASC """, (run_id,), ).fetchall() finalized_results: list[dict[str, object]] = [] activations: list[dict[str, object]] = [] for snapshot_row in snapshot_rows: document_id = int(snapshot_row["document_id"]) existing_result_row = find_active_result_row( connection, document_id=document_id, job_version_id=int(job_version_row["id"]), input_identity=str(snapshot_row["pinned_input_identity"]), ) if existing_result_row is not None and existing_result_row["created_text_revision_id"] is not None: result_id = int(existing_result_row["id"]) connection.execute( """ UPDATE run_items SET result_id = ? WHERE run_id = ? AND document_id = ? """, (result_id, run_id, document_id), ) activation_payload = maybe_activate_created_text_revision( connection, paths, run_row=run_row, job_version_row=job_version_row, document_id=document_id, result_id=result_id, text_revision_id=int(existing_result_row["created_text_revision_id"]), ) if activation_payload is not None: activations.append(activation_payload) finalized_results.append(result_summary_by_id(connection, result_id)) continue page_rows = connection.execute( """ SELECT * FROM run_items WHERE run_id = ? AND document_id = ? AND item_kind = 'page' ORDER BY page_number ASC, id ASC """, (run_id, document_id), ).fetchall() if not page_rows: raise RetrieverError( f"Image-description run {run_id} has no page items for document {document_id}." ) page_output_payloads = list_image_description_page_output_payloads_for_document( connection, run_id=run_id, document_id=document_id, ) if len(page_output_payloads) != len(page_rows): raise RetrieverError( f"Image-description run {run_id} is missing page outputs for document {document_id}: " f"expected {len(page_rows)}, found {len(page_output_payloads)}." ) merged_text = build_image_description_revision_text(page_output_payloads) quality_summary = text_quality_summary( merged_text, revision_kind="image_description", ) created_text_revision_id = create_text_revision_row( connection, paths, document_id=document_id, revision_kind="image_description", text_content=merged_text, language=None, parent_revision_id=( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row["pinned_input_revision_id"] is not None else None ), created_by_job_version_id=int(job_version_row["id"]), quality_score=float(quality_summary["quality_score"]), provider_metadata={ "run_id": run_id, "page_count": len(page_output_payloads), "finalized_from": "image_description_page_outputs", "quality": quality_summary, }, ) result_id, _ = create_result_row( connection, run_id=run_id, document_id=document_id, job_version_id=int(job_version_row["id"]), input_revision_id=( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row["pinned_input_revision_id"] is not None else None ), input_identity=str(snapshot_row["pinned_input_identity"]), raw_output={ "page_count": len(page_output_payloads), "finalized_from": "image_description_page_outputs", }, normalized_output={ "page_count": len(page_output_payloads), "created_text_revision_id": created_text_revision_id, }, created_text_revision_id=created_text_revision_id, provider_metadata={ "run_id": run_id, "page_count": len(page_output_payloads), "quality": quality_summary, }, ) connection.execute( """ UPDATE run_items SET result_id = ? WHERE run_id = ? AND document_id = ? """, (result_id, run_id, document_id), ) activation_payload = maybe_activate_created_text_revision( connection, paths, run_row=run_row, job_version_row=job_version_row, document_id=document_id, result_id=result_id, text_revision_id=created_text_revision_id, ) if activation_payload is not None: activations.append(activation_payload) finalized_results.append(result_summary_by_id(connection, result_id)) refresh_run_progress(connection, run_id) return { "status": "ok", "run": run_status_by_id(connection, run_id), "results": finalized_results, "activations": activations, } def coerce_result_output_value_for_publication(field_type: str, output_value: object) -> object: if output_value is None: return None if field_type == "text": return result_output_display_value(output_value) if isinstance(output_value, bool): raw_value = "true" if output_value else "false" elif isinstance(output_value, (int, float)): raw_value = str(output_value) elif isinstance(output_value, str): raw_value = output_value else: raw_value = compact_json_text(output_value) return value_from_type(field_type, raw_value) def publish_result_outputs_for_run( connection: sqlite3.Connection, *, run_id: int, output_names: list[str] | None = None, ) -> dict[str, object]: normalized_output_names = { sanitize_processing_identifier(raw_name, label="Job output name", prefix="output") for raw_name in (output_names or []) } result_output_rows = connection.execute( """ SELECT DISTINCT ro.id AS result_output_id, ro.result_id, ro.output_value_json, jo.id AS job_output_id, jo.output_name, jo.bound_custom_field, ri.document_id FROM run_items ri JOIN result_outputs ro ON ro.result_id = ri.result_id JOIN job_outputs jo ON jo.id = ro.job_output_id WHERE ri.run_id = ? AND ri.result_id IS NOT NULL AND jo.bound_custom_field IS NOT NULL ORDER BY ri.document_id ASC, jo.output_name ASC, ro.id ASC """, (run_id,), ).fetchall() published = 0 affected_documents: set[int] = set() published_outputs: list[dict[str, object]] = [] now = utc_now() for row in result_output_rows: output_name = str(row["output_name"]) if normalized_output_names and output_name not in normalized_output_names: continue custom_field_name = str(row["bound_custom_field"] or "") field_def = resolve_field_definition(connection, custom_field_name) if field_def["source"] != "custom": raise RetrieverError( f"Bound field {custom_field_name!r} for output {output_name!r} must be a custom field." ) document_id = int(row["document_id"]) output_value = decode_json_text(row["output_value_json"]) typed_value = coerce_result_output_value_for_publication(field_def["field_type"], output_value) connection.execute( f""" UPDATE documents SET {quote_identifier(field_def['field_name'])} = ?, updated_at = ? WHERE id = ? """, (typed_value, now, document_id), ) connection.execute( """ INSERT INTO publications ( result_output_id, document_id, job_output_id, custom_field_name, published_at ) VALUES (?, ?, ?, ?, ?) """, ( int(row["result_output_id"]), document_id, int(row["job_output_id"]), field_def["field_name"], now, ), ) published += 1 affected_documents.add(document_id) published_outputs.append( { "document_id": document_id, "result_id": int(row["result_id"]), "output_name": output_name, "custom_field_name": field_def["field_name"], "value": typed_value, } ) return { "published_count": published, "document_count": len(affected_documents), "published_outputs": published_outputs, } def run_summary_by_id(connection: sqlite3.Connection, run_id: int) -> dict[str, object]: row = connection.execute( """ SELECT * FROM runs WHERE id = ? """, (run_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown run id: {run_id}") payload = run_row_to_payload(row) payload["job_version"] = job_version_summary_by_id(connection, int(row["job_version_id"])) snapshot_rows = connection.execute( """ SELECT * FROM run_snapshot_documents WHERE run_id = ? ORDER BY ordinal ASC, id ASC """, (run_id,), ).fetchall() payload["documents"] = [run_snapshot_document_row_to_payload(snapshot_row) for snapshot_row in snapshot_rows] return payload def list_run_summaries(connection: sqlite3.Connection) -> list[dict[str, object]]: rows = connection.execute( """ SELECT * FROM runs ORDER BY id DESC """ ).fetchall() return [run_summary_by_id(connection, int(row["id"])) for row in rows] def resolve_from_run_snapshot_rows(connection: sqlite3.Connection, from_run_id: int) -> list[sqlite3.Row]: run_row = connection.execute("SELECT id FROM runs WHERE id = ?", (from_run_id,)).fetchone() if run_row is None: raise RetrieverError(f"Unknown run id: {from_run_id}") return connection.execute( """ SELECT * FROM run_snapshot_documents WHERE run_id = ? ORDER BY ordinal ASC, id ASC """, (from_run_id,), ).fetchall() def expand_seed_documents_with_family( connection: sqlite3.Connection, seed_document_ids: list[int], reasons_by_document_id: dict[int, dict[str, object]], ) -> list[int]: final_document_ids = list(seed_document_ids) seen_document_ids = set(seed_document_ids) if not seed_document_ids: return final_document_ids rows = connection.execute( f""" SELECT * FROM documents WHERE id IN ({', '.join('?' for _ in seed_document_ids)}) """, seed_document_ids, ).fetchall() row_by_id = {int(row["id"]): row for row in rows} for seed_document_id in seed_document_ids: seed_row = row_by_id.get(seed_document_id) if seed_row is None: continue family_root_id = int(seed_row["parent_document_id"]) if seed_row["parent_document_id"] is not None else int(seed_row["id"]) family_rows = connection.execute( """ SELECT * FROM documents WHERE (id = ? OR parent_document_id = ?) AND lifecycle_status NOT IN ('missing', 'deleted') ORDER BY CASE WHEN id = ? THEN 0 ELSE 1 END ASC, id ASC """, (family_root_id, family_root_id, family_root_id), ).fetchall() for family_row in family_rows: family_document_id = int(family_row["id"]) if family_document_id not in seen_document_ids: seen_document_ids.add(family_document_id) final_document_ids.append(family_document_id) reason_state = reasons_by_document_id.setdefault( family_document_id, {"direct_reasons": [], "family_seed_document_ids": []}, ) family_seed_ids = reason_state["family_seed_document_ids"] if seed_document_id not in family_seed_ids: family_seed_ids.append(seed_document_id) return final_document_ids def compute_source_parts_content_hash( connection: sqlite3.Connection, *, root: Path, document_id: int, ) -> str | None: source_part_rows = connection.execute( """ SELECT part_kind, rel_source_path, ordinal, label FROM document_source_parts WHERE document_id = ? ORDER BY ordinal ASC, id ASC """, (document_id,), ).fetchall() if not source_part_rows: return None signature_parts: list[dict[str, object]] = [] for row in source_part_rows: rel_source_path = normalize_whitespace(str(row["rel_source_path"] or "")) absolute_path = (root / rel_source_path).resolve() if rel_source_path else None file_hash = sha256_file(absolute_path) if absolute_path is not None and absolute_path.exists() else None signature_parts.append( { "part_kind": normalize_whitespace(str(row["part_kind"] or "")), "rel_source_path": rel_source_path, "ordinal": int(row["ordinal"] or 0), "label": normalize_whitespace(str(row["label"] or "")) or None, "file_hash": file_hash, } ) return sha256_text(compact_json_text(signature_parts)) def compute_document_input_reference_for_job_version( connection: sqlite3.Connection, *, root: Path, document_row: sqlite3.Row, job_row: sqlite3.Row, job_version_row: sqlite3.Row, frozen_input_revision_id: int | None, frozen_content_hash: str | None, ) -> dict[str, object]: input_basis = str(job_version_row["input_basis"]) job_kind = str(job_row["job_kind"]) capability = normalize_whitespace(str(job_version_row["capability"] or "")) provider = normalize_whitespace(str(job_version_row["provider"] or "")) model = normalize_whitespace(str(job_version_row["model"] or "")) backend_id = capability or (f"{provider}:{model}" if model else provider) parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} if input_basis in {"active_search_text", "source_extract", "text_revision"}: if frozen_input_revision_id is not None: revision_id = int(frozen_input_revision_id) elif input_basis == "source_extract": revision_id = ( int(document_row["source_text_revision_id"]) if document_row["source_text_revision_id"] is not None else None ) else: revision_id = ( int(document_row["active_search_text_revision_id"]) if document_row["active_search_text_revision_id"] is not None else None ) if revision_id is None: raise RetrieverError( f"Document {document_row['id']} has no pinned text revision for input basis {input_basis!r}. Reingest the document first." ) revision_row = connection.execute( """ SELECT content_hash FROM text_revisions WHERE id = ? AND retracted_at IS NULL """, (revision_id,), ).fetchone() if revision_row is None: raise RetrieverError(f"Unknown active text revision id: {revision_id}") pinned_content_hash = frozen_content_hash or str(revision_row["content_hash"]) if job_kind == "translation": target_language = ( parameters.get("target_language") or parameters.get("target_lang") or parameters.get("language") ) if not isinstance(target_language, str) or not target_language.strip(): raise RetrieverError("Translation job versions require parameters_json.target_language.") input_identity = build_translation_input_identity(revision_id, target_language=target_language) else: input_identity = build_text_revision_input_identity(revision_id) return { "pinned_input_revision_id": revision_id, "pinned_input_identity": input_identity, "pinned_content_hash": pinned_content_hash, } if input_basis in {"source_file", "source_parts"}: source_hash = normalize_whitespace(str(frozen_content_hash or document_row["file_hash"] or "")) if not source_hash and input_basis == "source_parts": source_hash = normalize_whitespace( str(compute_source_parts_content_hash(connection, root=root, document_id=int(document_row["id"])) or "") ) if not source_hash: raise RetrieverError( f"Document {document_row['id']} has no file hash for input basis {input_basis!r}." ) if job_kind == "ocr": rendering_settings = parameters.get("rendering_settings") normalized_settings = rendering_settings if isinstance(rendering_settings, dict) else {} input_identity = build_ocr_input_identity( source_hash, rendering_settings=normalized_settings, backend_id=backend_id, ) else: image_prep_settings = parameters.get("image_prep_settings") normalized_settings = image_prep_settings if isinstance(image_prep_settings, dict) else {} input_identity = build_image_source_input_identity( source_hash, image_prep_settings=normalized_settings, backend_id=backend_id, ) return { "pinned_input_revision_id": None, "pinned_input_identity": input_identity, "pinned_content_hash": source_hash, } raise RetrieverError(f"Unsupported job input basis for run planning: {input_basis}") def plan_scope_run_snapshot_rows( connection: sqlite3.Connection, *, root: Path, job_row: sqlite3.Row, job_version_row: sqlite3.Row, selector: dict[str, object], family_mode: str, seed_limit: int | None, ) -> list[dict[str, object]]: selected_documents, frozen_inputs_by_document_id = plan_scope_selected_documents( connection, selector=selector, family_mode=family_mode, seed_limit=seed_limit, ) if not selected_documents: return [] snapshot_rows: list[dict[str, object]] = [] for selected_document in selected_documents: document_id = int(selected_document["document_id"]) document_row = selected_document["document_row"] frozen_inputs = frozen_inputs_by_document_id.get(document_id, {}) pinned_input = compute_document_input_reference_for_job_version( connection, root=root, document_row=document_row, job_row=job_row, job_version_row=job_version_row, frozen_input_revision_id=frozen_inputs.get("pinned_input_revision_id"), frozen_content_hash=frozen_inputs.get("pinned_content_hash"), ) snapshot_rows.append( { "document_id": document_id, "ordinal": int(selected_document["ordinal"]), "inclusion_reason": selected_document["inclusion_reason"], "pinned_input_revision_id": pinned_input["pinned_input_revision_id"], "pinned_input_identity": pinned_input["pinned_input_identity"], "pinned_content_hash": pinned_input["pinned_content_hash"], } ) return snapshot_rows def scope_run_selector_has_inputs(selector: dict[str, object]) -> bool: return bool(scope_selector_instances(selector)) def scope_selector_reason_entries(scope: dict[str, object]) -> list[dict[str, object]]: reasons: list[dict[str, object]] = [] keyword = normalize_inline_whitespace(str(scope.get("keyword") or "")) if keyword: reasons.append({"type": "keyword", "query": keyword}) bates = scope.get("bates") if isinstance(bates, dict): begin = normalize_inline_whitespace(str(bates.get("begin") or "")) end = normalize_inline_whitespace(str(bates.get("end") or "")) if begin and end: reasons.append({"type": "bates", "begin": begin, "end": end}) filter_expression = normalize_inline_whitespace(str(scope.get("filter") or "")) if filter_expression: reasons.append({"type": "filter", "expression": filter_expression}) dataset_entries = coerce_scope_dataset_entries(scope.get("dataset")) if dataset_entries: reasons.append( { "type": "dataset", "dataset_ids": [int(entry["id"]) for entry in dataset_entries], "dataset_names": [str(entry["name"]) for entry in dataset_entries], } ) if scope.get("from_run_id") is not None: reasons.append({"type": "from_run", "run_id": int(scope["from_run_id"])}) return reasons def resolve_seed_documents_for_scope_selector( connection: sqlite3.Connection, selector: dict[str, object], ) -> tuple[list[int], dict[int, dict[str, object]], dict[int, dict[str, object]]]: selector_instances = scope_selector_instances(selector) if not selector_instances: return [], {}, {} reasons_by_document_id: dict[int, dict[str, object]] = {} frozen_inputs_by_document_id: dict[int, dict[str, object]] = {} ordered_results: list[dict[str, object]] | None = None matching_document_ids: set[int] | None = None for selector_instance in selector_instances: selection = resolve_scope_document_search(connection, selector_instance) selection_scope = coerce_scope_payload(selection.get("scope")) selection_results = selection["results"] selection_document_ids = {int(item["id"]) for item in selection_results} if ordered_results is None: ordered_results = selection_results matching_document_ids = selection_document_ids else: assert matching_document_ids is not None matching_document_ids &= selection_document_ids direct_reasons = scope_selector_reason_entries(selection_scope) if direct_reasons: for item in selection_results: document_id = int(item["id"]) reason_state = reasons_by_document_id.setdefault( document_id, {"direct_reasons": [], "family_seed_document_ids": []}, ) for reason in direct_reasons: if reason not in reason_state["direct_reasons"]: reason_state["direct_reasons"].append(reason) from_run_id = selection_scope.get("from_run_id") if from_run_id is not None: # Later selector clauses win so explicit --from-run-id narrows and pins over scope. for snapshot_row in resolve_from_run_snapshot_rows(connection, int(from_run_id)): document_id = int(snapshot_row["document_id"]) frozen_inputs_by_document_id[document_id] = { "pinned_input_revision_id": ( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row["pinned_input_revision_id"] is not None else None ), "pinned_content_hash": snapshot_row["pinned_content_hash"], } if ordered_results is None or matching_document_ids is None: return [], {}, {} ordered_matching_ids = [ int(item["id"]) for item in ordered_results if int(item["id"]) in matching_document_ids ] filtered_reasons = { document_id: reasons_by_document_id.get( document_id, {"direct_reasons": [], "family_seed_document_ids": []}, ) for document_id in ordered_matching_ids } filtered_frozen_inputs = { document_id: frozen_inputs for document_id, frozen_inputs in frozen_inputs_by_document_id.items() if document_id in matching_document_ids } return ordered_matching_ids, filtered_reasons, filtered_frozen_inputs def plan_scope_selected_documents( connection: sqlite3.Connection, *, selector: dict[str, object], family_mode: str, seed_limit: int | None, ) -> tuple[list[dict[str, object]], dict[int, dict[str, object]]]: seed_document_ids, reasons_by_document_id, frozen_inputs_by_document_id = resolve_seed_documents_for_scope_selector( connection, selector, ) if not seed_document_ids: return [], frozen_inputs_by_document_id if seed_limit is not None: seed_document_ids = seed_document_ids[:seed_limit] if family_mode == "with_family": final_document_ids = expand_seed_documents_with_family(connection, seed_document_ids, reasons_by_document_id) else: final_document_ids = list(seed_document_ids) if not final_document_ids: return [], frozen_inputs_by_document_id document_rows = connection.execute( f""" SELECT * FROM documents WHERE id IN ({', '.join('?' for _ in final_document_ids)}) ORDER BY id ASC """, final_document_ids, ).fetchall() document_row_by_id = {int(row["id"]): row for row in document_rows} selected_documents: list[dict[str, object]] = [] for ordinal, document_id in enumerate(final_document_ids): document_row = document_row_by_id.get(int(document_id)) if document_row is None: continue selected_documents.append( { "document_id": int(document_id), "ordinal": ordinal, "document_row": document_row, "inclusion_reason": reasons_by_document_id.get( int(document_id), {"direct_reasons": [], "family_seed_document_ids": []}, ), } ) return selected_documents, frozen_inputs_by_document_id def load_run_snapshot_rows(connection: sqlite3.Connection, run_id: int) -> list[sqlite3.Row]: return connection.execute( """ SELECT * FROM run_snapshot_documents WHERE run_id = ? ORDER BY ordinal ASC, id ASC """, (run_id,), ).fetchall() def load_text_input_for_snapshot_row( connection: sqlite3.Connection, paths: dict[str, Path], snapshot_row: sqlite3.Row, ) -> str | None: pinned_input_revision_id = snapshot_row["pinned_input_revision_id"] if pinned_input_revision_id is None: return None revision_row = require_text_revision_row_by_id(connection, int(pinned_input_revision_id)) revision_body = read_text_revision_body(paths, revision_row["storage_rel_path"]) if revision_body is None: raise RetrieverError(f"Text revision {pinned_input_revision_id} has no readable body on disk.") return revision_body def job_output_schema_fragment(value_type: str) -> dict[str, object]: normalized_value_type = normalize_job_output_value_type(value_type) if normalized_value_type == "boolean": return {"type": "boolean"} if normalized_value_type == "date": return {"type": "string", "description": "ISO-8601 date string"} if normalized_value_type == "integer": return {"type": "integer"} if normalized_value_type == "json": return {"type": "object"} if normalized_value_type == "real": return {"type": "number"} return {"type": "string"} def derive_job_output_response_schema(job_output_rows: list[sqlite3.Row]) -> dict[str, object]: properties: dict[str, object] = {} required: list[str] = [] for row in job_output_rows: output_name = str(row["output_name"]) properties[output_name] = job_output_schema_fragment(str(row["value_type"])) required.append(output_name) return { "type": "object", "properties": properties, "required": required, "additionalProperties": False, } def default_job_output_value(value_type: str) -> object: normalized_value_type = normalize_job_output_value_type(value_type) if normalized_value_type == "boolean": return False if normalized_value_type == "integer": return 0 if normalized_value_type == "real": return 0 if normalized_value_type == "json": return {} return "" def build_text_structured_execution_payload( job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], ) -> dict[str, object]: output_defaults = { str(row["output_name"]): default_job_output_value(str(row["value_type"])) for row in job_output_rows } instruction_text = normalize_whitespace(str(job_version_row["instruction_text"] or "")) prompt_lines = [ "Read the provided document text and extract the requested outputs.", "Return only a JSON object that matches response_schema exactly.", "If a value cannot be determined, use the matching fallback value from output_defaults.", ] if instruction_text: prompt_lines.append(f"Job instruction: {instruction_text}") return { "capability": "text_structured", "task_prompt": "\n".join(prompt_lines), "output_defaults": output_defaults, "completion_command": "complete-run-item", "completion_template": { "raw_output_json": {key: f"" for key in output_defaults}, "normalized_output_json": {key: f"" for key in output_defaults}, "output_values_json": {key: f"" for key in output_defaults}, }, } def build_translation_execution_payload(job_version_row: sqlite3.Row) -> dict[str, object]: parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} target_language = normalize_whitespace( str( parameters.get("target_language") or parameters.get("target_lang") or parameters.get("language") or "" ) ) instruction_text = normalize_whitespace(str(job_version_row["instruction_text"] or "")) prompt_lines = [ f"Translate the entire input text into {target_language or 'the requested target language'}.", "Return only the translated text, with no JSON wrapper or commentary.", "Preserve the original structure and line breaks.", "Keep dates, numbers, email addresses, URLs, Bates numbers, file names, and header labels/order unchanged.", "Keep proper names unchanged unless the source itself provides a translated form.", "Do not summarize or omit content.", ] if instruction_text: prompt_lines.append(f"Job instruction: {instruction_text}") return { "capability": "text_translation", "target_language": target_language or None, "task_prompt": "\n".join(prompt_lines), "completion_command": "complete-run-item", "completion_template": { "raw_output_json": { "translated_text": "", "target_language": target_language or "", }, "normalized_output_json": { "translated_text": "", "target_language": target_language or "", }, "created_text_revision_json": { "revision_kind": "translation", "language": target_language or "", "text_content": "", }, }, } def build_vision_ocr_execution_payload(job_version_row: sqlite3.Row) -> dict[str, object]: instruction_text = normalize_whitespace(str(job_version_row["instruction_text"] or "")) prompt_lines = [ "Read the page image and transcribe all readable text in natural reading order.", "Use the original source image referenced by input.artifact_path whenever possible.", "Do not rely on resized previews, screenshots, or re-encoded copies unless the original artifact is unavailable or unreadable.", "Return plain text only. Do not wrap the output in JSON or Markdown.", ] if instruction_text: prompt_lines.append(f"Job instruction: {instruction_text}") return { "capability": "vision_ocr", "task_prompt": "\n".join(prompt_lines), "completion_command": "complete-run-item", "completion_template": { "page_text": "", }, } def build_vision_description_execution_payload(job_version_row: sqlite3.Row) -> dict[str, object]: instruction_text = normalize_whitespace(str(job_version_row["instruction_text"] or "")) prompt_lines = [ "Describe the page image in clear, search-friendly prose.", "Prioritize visible people, objects, scenes, charts, handwriting, stamps, and other salient visual details.", "Return plain text only. Do not wrap the output in JSON or Markdown.", ] if instruction_text: prompt_lines.append(f"Job instruction: {instruction_text}") return { "capability": "vision_description", "task_prompt": "\n".join(prompt_lines), "completion_command": "complete-run-item", "completion_template": { "page_text": "", }, } def build_capability_execution_payload( job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], ) -> dict[str, object] | None: capability = normalize_whitespace(str(job_version_row["capability"] or "")) if capability == "text_structured": return build_text_structured_execution_payload(job_version_row, job_output_rows) if capability == "text_translation": return build_translation_execution_payload(job_version_row) if capability == "vision_description": return build_vision_description_execution_payload(job_version_row) if capability == "vision_ocr": return build_vision_ocr_execution_payload(job_version_row) return None def best_effort_document_source_path(root: Path, document_row: sqlite3.Row) -> str | None: source_rel_path = normalize_whitespace(str(document_row["source_rel_path"] or "")) if not source_rel_path: return None candidate = (root / source_rel_path).resolve() if candidate.exists(): return str(candidate) return None def build_text_input_reference( paths: dict[str, Path], revision_row: sqlite3.Row, *, inline_bytes: int = DEFAULT_RUN_ITEM_CONTEXT_INLINE_BYTES, ) -> dict[str, object]: revision_body = read_text_revision_body(paths, revision_row["storage_rel_path"]) if revision_body is None: raise RetrieverError(f"Text revision {revision_row['id']} has no readable body on disk.") encoded_size = len(revision_body.encode("utf-8")) storage_path = str((paths["state_dir"] / str(revision_row["storage_rel_path"])).resolve()) return { "kind": "text_revision", "text_revision_id": int(revision_row["id"]), "inline_text": revision_body if encoded_size <= inline_bytes else None, "text_path": None if encoded_size <= inline_bytes else storage_path, "bytes": encoded_size, } def build_run_item_context_payload( connection: sqlite3.Connection, paths: dict[str, Path], root: Path, run_item_row: sqlite3.Row, *, inline_bytes: int = DEFAULT_RUN_ITEM_CONTEXT_INLINE_BYTES, ) -> dict[str, object]: run_row = require_run_row_by_id(connection, int(run_item_row["run_id"])) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( """ SELECT * FROM jobs WHERE id = ? """, (job_version_row["job_id"],), ).fetchone() assert job_row is not None snapshot_row = None if run_item_row["run_snapshot_document_id"] is not None: snapshot_row = connection.execute( """ SELECT * FROM run_snapshot_documents WHERE id = ? """, (run_item_row["run_snapshot_document_id"],), ).fetchone() document_row = connection.execute( """ SELECT * FROM documents WHERE id = ? """, (run_item_row["document_id"],), ).fetchone() if document_row is None: raise RetrieverError(f"Unknown document id: {run_item_row['document_id']}") job_output_rows = connection.execute( """ SELECT * FROM job_outputs WHERE job_id = ? ORDER BY ordinal ASC, output_name ASC, id ASC """, (job_row["id"],), ).fetchall() response_schema = decode_json_text(job_version_row["response_schema_json"]) if response_schema is None and str(job_version_row["capability"]) == "text_structured": response_schema = derive_job_output_response_schema(job_output_rows) item_kind = str(run_item_row["item_kind"] or "") if item_kind == "page": artifact_rel_path = normalize_whitespace(str(run_item_row["input_artifact_rel_path"] or "")) artifact_path = resolve_workspace_artifact_path(root, artifact_rel_path) if artifact_path is None or not artifact_path.exists(): raise RetrieverError(f"Run item {run_item_row['id']} points at a missing OCR artifact: {artifact_rel_path!r}") read_artifact_path = ensure_read_safe_visual_artifact_path(root, artifact_path) try: read_artifact_rel_path = read_artifact_path.relative_to(root).as_posix() except ValueError: read_artifact_rel_path = None page_input_kind = ( "image_description_page_image" if normalize_job_kind(str(job_row["job_kind"])) == "image_description" else "ocr_page_image" ) input_payload = { "kind": page_input_kind, "page_number": int(run_item_row["page_number"] or 0), "artifact_rel_path": artifact_rel_path, "artifact_path": str(artifact_path), "read_artifact_rel_path": read_artifact_rel_path, "read_artifact_path": str(read_artifact_path), "bytes": artifact_path.stat().st_size, "text_revision_id": None, "inline_text": None, "text_path": None, "source_path": best_effort_document_source_path(root, document_row), "source_rel_path": document_row["source_rel_path"], "file_hash": document_row["file_hash"], } elif snapshot_row is not None and snapshot_row["pinned_input_revision_id"] is not None: revision_row = require_text_revision_row_by_id(connection, int(snapshot_row["pinned_input_revision_id"])) input_payload = build_text_input_reference(paths, revision_row, inline_bytes=inline_bytes) else: input_payload = { "kind": str(job_version_row["input_basis"]), "text_revision_id": None, "inline_text": None, "text_path": None, "bytes": None, "source_path": best_effort_document_source_path(root, document_row), "source_rel_path": document_row["source_rel_path"], "file_hash": document_row["file_hash"], } return { "run": run_row_to_payload(run_row), "run_item": run_item_row_to_payload(run_item_row), "job": job_row_to_payload(job_row), "job_version": job_version_row_to_payload(job_version_row), "job_outputs": [job_output_row_to_payload(row) for row in job_output_rows], "response_schema": response_schema, "execution": build_capability_execution_payload(job_version_row, job_output_rows), "document": { "id": int(document_row["id"]), "file_name": document_row["file_name"], "file_type": document_row["file_type"], "source_kind": document_row["source_kind"], "source_rel_path": document_row["source_rel_path"], "source_path": best_effort_document_source_path(root, document_row), "active_search_text_revision_id": document_row["active_search_text_revision_id"], "source_text_revision_id": document_row["source_text_revision_id"], }, "input": input_payload, "inclusion_reason": ( decode_json_text(snapshot_row["inclusion_reason_json"], default={}) or {} if snapshot_row is not None else {} ), } OPENAI_RESPONSES_PROVIDER_NAMES = {"openai", "openai_responses"} STATIC_STRUCTURED_EXTRACTION_PROVIDER_NAMES = {"builtin_static_json", "static_json"} STATIC_TRANSLATION_PROVIDER_NAMES = {"builtin_static_text", "static_text"} OPENAI_RESPONSES_DEFAULT_URL = "https://api.openai.com/v1/responses" def render_processing_template_value(value: object, context: dict[str, object]) -> object: if isinstance(value, str): normalized_context = defaultdict( str, { key: "" if item is None else str(item) for key, item in context.items() }, ) try: return value.format_map(normalized_context) except Exception: return value if isinstance(value, list): return [render_processing_template_value(item, context) for item in value] if isinstance(value, dict): return { str(key): render_processing_template_value(item, context) for key, item in value.items() } return value def processing_job_output_json_schema(job_output_rows: list[sqlite3.Row]) -> dict[str, object]: properties: dict[str, object] = {} required: list[str] = [] for job_output_row in job_output_rows: output_name = str(job_output_row["output_name"]) value_type = normalize_job_output_value_type(str(job_output_row["value_type"] or "text")) schema: dict[str, object] if value_type == "boolean": schema = {"type": "boolean"} elif value_type == "date": schema = {"type": "string"} elif value_type == "integer": schema = {"type": "integer"} elif value_type == "real": schema = {"type": "number"} elif value_type == "json": schema = { "type": "object", "additionalProperties": True, } else: schema = {"type": "string"} properties[output_name] = schema required.append(output_name) return { "type": "object", "properties": properties, "required": required, "additionalProperties": False, } def processing_response_schema( job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], ) -> dict[str, object]: parsed = decode_json_text(job_version_row["response_schema_json"]) if isinstance(parsed, dict) and parsed: return parsed return processing_job_output_json_schema(job_output_rows) def json_schema_path(path: str, segment: object) -> str: if isinstance(segment, int): return f"{path}[{segment}]" if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", str(segment or "")): return f"{path}.{segment}" return f'{path}[{json.dumps(str(segment))}]' def json_schema_allowed_types(schema: dict[str, object]) -> list[str]: raw_type = schema.get("type") if isinstance(raw_type, str): return [raw_type] if isinstance(raw_type, list): return [str(item) for item in raw_type if isinstance(item, str)] return [] def json_schema_type_matches(value: object, schema_type: str) -> bool: if schema_type == "object": return isinstance(value, dict) if schema_type == "array": return isinstance(value, list) if schema_type == "string": return isinstance(value, str) if schema_type == "integer": return isinstance(value, int) and not isinstance(value, bool) if schema_type == "number": return (isinstance(value, int) and not isinstance(value, bool)) or isinstance(value, float) if schema_type == "boolean": return isinstance(value, bool) if schema_type == "null": return value is None return True def validate_processing_schema_value( value: object, schema: dict[str, object], *, path: str = "$", ) -> list[str]: if not isinstance(schema, dict): return [] issues: list[str] = [] allowed_types = json_schema_allowed_types(schema) if allowed_types and not any(json_schema_type_matches(value, schema_type) for schema_type in allowed_types): issues.append(f"{path} must be of type {' or '.join(allowed_types)}") return issues enum_values = schema.get("enum") if isinstance(enum_values, list) and enum_values and value not in enum_values: issues.append(f"{path} must be one of {enum_values!r}") if isinstance(value, dict): properties = schema.get("properties") property_schemas = properties if isinstance(properties, dict) else {} required = schema.get("required") if isinstance(required, list): for property_name in required: if isinstance(property_name, str) and property_name not in value: issues.append(f"{json_schema_path(path, property_name)} is required") additional_properties = schema.get("additionalProperties", True) for key, item in value.items(): child_schema = property_schemas.get(key) if isinstance(child_schema, dict): issues.extend( validate_processing_schema_value( item, child_schema, path=json_schema_path(path, key), ) ) continue if additional_properties is False: issues.append(f"{json_schema_path(path, key)} is not allowed") elif isinstance(additional_properties, dict): issues.extend( validate_processing_schema_value( item, additional_properties, path=json_schema_path(path, key), ) ) return issues if isinstance(value, list): items_schema = schema.get("items") if isinstance(items_schema, dict): for index, item in enumerate(value): issues.extend( validate_processing_schema_value( item, items_schema, path=json_schema_path(path, index), ) ) return issues return issues def validate_processing_response_output( job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], output_value: object, ) -> list[str]: schema = processing_response_schema(job_version_row, job_output_rows) return validate_processing_schema_value(output_value, schema) def openai_responses_api_url() -> str: explicit_url = normalize_whitespace(os.environ.get("OPENAI_RESPONSES_URL", "")) if explicit_url: return explicit_url base_url = normalize_whitespace(os.environ.get("OPENAI_BASE_URL", "")) if base_url: return base_url.rstrip("/") + "/responses" return OPENAI_RESPONSES_DEFAULT_URL def openai_api_key() -> str: api_key = normalize_whitespace(os.environ.get("OPENAI_API_KEY", "")) if not api_key: raise RetrieverError("OPENAI_API_KEY is required for openai_responses providers.") return api_key def openai_response_text(response_payload: dict[str, object]) -> str: top_level_output_text = response_payload.get("output_text") if isinstance(top_level_output_text, str) and top_level_output_text.strip(): return top_level_output_text output_items = response_payload.get("output") if not isinstance(output_items, list): raise RetrieverError("OpenAI response payload did not contain output text.") text_parts: list[str] = [] for output_item in output_items: if not isinstance(output_item, dict): continue content_items = output_item.get("content") if not isinstance(content_items, list): continue for content_item in content_items: if not isinstance(content_item, dict): continue if str(content_item.get("type") or "") != "output_text": continue text_value = content_item.get("text") if isinstance(text_value, str): text_parts.append(text_value) if text_parts: return "".join(text_parts) raise RetrieverError("OpenAI response payload did not contain output_text content.") def openai_response_usage(response_payload: dict[str, object]) -> tuple[int | None, int | None]: usage = response_payload.get("usage") if not isinstance(usage, dict): return None, None input_tokens = usage.get("input_tokens") output_tokens = usage.get("output_tokens") return ( int(input_tokens) if isinstance(input_tokens, int) else None, int(output_tokens) if isinstance(output_tokens, int) else None, ) def call_openai_responses_api( *, payload: dict[str, object], timeout_seconds: float, ) -> dict[str, object]: request_bytes = compact_json_text(payload).encode("utf-8") request = urllib_request.Request( openai_responses_api_url(), data=request_bytes, headers={ "Authorization": f"Bearer {openai_api_key()}", "Content-Type": "application/json", }, method="POST", ) try: with urllib_request.urlopen(request, timeout=timeout_seconds) as response: response_body = response.read().decode("utf-8") except urllib_error.HTTPError as exc: error_body = exc.read().decode("utf-8", errors="replace") raise RetrieverError(f"OpenAI Responses API request failed: HTTP {exc.code}: {error_body}") from exc except urllib_error.URLError as exc: raise RetrieverError(f"OpenAI Responses API request failed: {exc.reason}") from exc try: parsed = json.loads(response_body) except (TypeError, ValueError, json.JSONDecodeError) as exc: raise RetrieverError("OpenAI Responses API returned invalid JSON.") from exc if not isinstance(parsed, dict): raise RetrieverError("OpenAI Responses API returned an unexpected response shape.") return parsed def response_api_request_overrides(parameters: dict[str, object]) -> dict[str, object]: overrides: dict[str, object] = {} for key in ( "max_output_tokens", "metadata", "reasoning", "service_tier", "temperature", "top_p", "truncation", ): if key in parameters: overrides[key] = parameters[key] return overrides def processing_input_context(document_row: sqlite3.Row, text_input: str | None) -> dict[str, object]: return { "document_id": int(document_row["id"]), "control_number": document_row["control_number"], "file_name": document_row["file_name"], "rel_path": document_row["rel_path"], "title": document_row["title"], "subject": document_row["subject"], "text": text_input or "", } def execute_static_structured_extraction_provider( *, job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], document_row: sqlite3.Row, text_input: str | None, ) -> dict[str, object]: parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} output_values = parameters.get("output_values") if not isinstance(output_values, dict): raise RetrieverError("static_json structured extraction requires parameters_json.output_values to be an object.") context = processing_input_context(document_row, text_input) normalized_outputs: dict[str, object] = {} for job_output_row in job_output_rows: output_name = str(job_output_row["output_name"]) normalized_outputs[output_name] = render_processing_template_value(output_values.get(output_name), context) output_json = compact_json_text(normalized_outputs) return { "raw_output": normalized_outputs, "normalized_output": normalized_outputs, "output_values": normalized_outputs, "provider_request_id": None, "input_tokens": token_estimate(text_input or "") if text_input is not None else None, "output_tokens": token_estimate(output_json) if normalized_outputs else 0, "cost_cents": 0, "provider_metadata": { "provider": str(job_version_row["provider"]), "model": job_version_row["model"], "executed_by": "structured_extraction_static_provider", }, } def build_openai_structured_extraction_payload( *, job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], document_row: sqlite3.Row, text_input: str | None, ) -> dict[str, object]: model = normalize_whitespace(str(job_version_row["model"] or "")) if not model: raise RetrieverError("openai_responses structured extraction requires a model name.") parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} instruction_text = str(job_version_row["instruction_text"] or "").strip() or "Extract the requested structured outputs." schema = processing_response_schema(job_version_row, job_output_rows) schema_name = sanitize_processing_identifier( str(job_version_row["display_name"] or f"job_version_{job_version_row['id']}"), label="Schema name", prefix="schema", )[:64] user_content = ( "Return structured JSON for the following document text.\n\n" f"\n{text_input or ''}\n" ) payload: dict[str, object] = { "model": model, "store": False, "input": [ {"role": "system", "content": instruction_text}, {"role": "user", "content": user_content}, ], "text": { "format": { "type": "json_schema", "name": schema_name, "schema": schema, "strict": True, } }, } payload.update(response_api_request_overrides(parameters)) return payload def execute_openai_structured_extraction_provider( *, job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], document_row: sqlite3.Row, text_input: str | None, ) -> dict[str, object]: parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} timeout_seconds = float(parameters.get("timeout_seconds") or 60.0) request_payload = build_openai_structured_extraction_payload( job_version_row=job_version_row, job_output_rows=job_output_rows, document_row=document_row, text_input=text_input, ) response_payload = call_openai_responses_api( payload=request_payload, timeout_seconds=timeout_seconds, ) output_text = openai_response_text(response_payload) try: normalized_outputs = json.loads(output_text) except (TypeError, ValueError, json.JSONDecodeError) as exc: raise RetrieverError("OpenAI structured extraction did not return valid JSON.") from exc if not isinstance(normalized_outputs, dict): raise RetrieverError("OpenAI structured extraction must return a JSON object.") input_tokens, output_tokens = openai_response_usage(response_payload) return { "raw_output": response_payload, "normalized_output": normalized_outputs, "output_values": normalized_outputs, "provider_request_id": response_payload.get("id"), "input_tokens": input_tokens, "output_tokens": output_tokens, "cost_cents": None, "provider_metadata": { "provider": str(job_version_row["provider"]), "model": job_version_row["model"], "response_id": response_payload.get("id"), "status": response_payload.get("status"), }, } def execute_static_translation_provider( *, job_version_row: sqlite3.Row, document_row: sqlite3.Row, text_input: str | None, ) -> dict[str, object]: parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} translated_template = parameters.get("translated_text") if translated_template is None: raise RetrieverError("static_text translation requires parameters_json.translated_text.") context = processing_input_context(document_row, text_input) translated_text = render_processing_template_value(translated_template, context) if not isinstance(translated_text, str): translated_text = compact_json_text(translated_text) target_language = normalize_whitespace( str(parameters.get("target_language") or parameters.get("target_lang") or parameters.get("language") or "") ).lower() or None return { "raw_output": {"translated_text": translated_text}, "normalized_output": {"translated_text": translated_text}, "output_values": {}, "created_text_revision": { "revision_kind": "translation", "text_content": translated_text, "language": target_language, }, "provider_request_id": None, "input_tokens": token_estimate(text_input or "") if text_input is not None else None, "output_tokens": token_estimate(translated_text), "cost_cents": 0, "provider_metadata": { "provider": str(job_version_row["provider"]), "model": job_version_row["model"], "executed_by": "translation_static_provider", "target_language": target_language, }, } def build_openai_translation_payload( *, job_version_row: sqlite3.Row, text_input: str | None, ) -> tuple[dict[str, object], str | None]: model = normalize_whitespace(str(job_version_row["model"] or "")) if not model: raise RetrieverError("openai_responses translation requires a model name.") parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} target_language = normalize_whitespace( str(parameters.get("target_language") or parameters.get("target_lang") or parameters.get("language") or "") ).lower() if not target_language: raise RetrieverError("Translation job versions require parameters_json.target_language.") instruction_text = str(job_version_row["instruction_text"] or "").strip() prompt_lines = [ f"Translate the document into {target_language}. Return only the translated text.", "Preserve the original structure and line breaks.", "Keep dates, numbers, email addresses, URLs, Bates numbers, file names, and header labels/ordering unchanged.", "Keep proper names unchanged unless the source itself provides a translated form.", "Do not summarize or omit content.", ] if instruction_text: prompt_lines.append(f"Job instruction: {instruction_text}") payload: dict[str, object] = { "model": model, "store": False, "input": [ {"role": "system", "content": "\n".join(prompt_lines)}, { "role": "user", "content": f"\n{text_input or ''}\n", }, ], "text": {"format": {"type": "text"}}, } payload.update(response_api_request_overrides(parameters)) return payload, target_language def execute_openai_translation_provider( *, job_version_row: sqlite3.Row, text_input: str | None, ) -> dict[str, object]: parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(parameters, dict): parameters = {} timeout_seconds = float(parameters.get("timeout_seconds") or 60.0) request_payload, target_language = build_openai_translation_payload( job_version_row=job_version_row, text_input=text_input, ) response_payload = call_openai_responses_api( payload=request_payload, timeout_seconds=timeout_seconds, ) translated_text = openai_response_text(response_payload) input_tokens, output_tokens = openai_response_usage(response_payload) return { "raw_output": response_payload, "normalized_output": {"translated_text": translated_text}, "output_values": {}, "created_text_revision": { "revision_kind": "translation", "text_content": translated_text, "language": target_language, }, "provider_request_id": response_payload.get("id"), "input_tokens": input_tokens, "output_tokens": output_tokens, "cost_cents": None, "provider_metadata": { "provider": str(job_version_row["provider"]), "model": job_version_row["model"], "response_id": response_payload.get("id"), "status": response_payload.get("status"), "target_language": target_language, }, } async def execute_job_provider( *, job_row: sqlite3.Row, job_version_row: sqlite3.Row, job_output_rows: list[sqlite3.Row], document_row: sqlite3.Row, text_input: str | None, ) -> dict[str, object]: provider = normalize_whitespace(str(job_version_row["provider"] or "")).lower() job_kind = normalize_job_kind(str(job_row["job_kind"])) if job_kind == "structured_extraction": if provider in STATIC_STRUCTURED_EXTRACTION_PROVIDER_NAMES: return execute_static_structured_extraction_provider( job_version_row=job_version_row, job_output_rows=job_output_rows, document_row=document_row, text_input=text_input, ) if provider in OPENAI_RESPONSES_PROVIDER_NAMES: return await asyncio.to_thread( execute_openai_structured_extraction_provider, job_version_row=job_version_row, job_output_rows=job_output_rows, document_row=document_row, text_input=text_input, ) if job_kind == "translation": if provider in STATIC_TRANSLATION_PROVIDER_NAMES: return execute_static_translation_provider( job_version_row=job_version_row, document_row=document_row, text_input=text_input, ) if provider in OPENAI_RESPONSES_PROVIDER_NAMES: return await asyncio.to_thread( execute_openai_translation_provider, job_version_row=job_version_row, text_input=text_input, ) raise RetrieverError( f"Unsupported job execution provider {provider!r} for job kind {job_kind!r}." ) CONTAINER_SOURCE_FILE_TYPES = frozenset({PST_SOURCE_KIND, MBOX_SOURCE_KIND}) def default_ingest_stats(slack_export_count: int, gmail_export_count: int) -> dict[str, int]: return { "new": 0, "updated": 0, "renamed": 0, "missing": 0, "skipped": 0, "failed": 0, "pst_sources_skipped": 0, "pst_messages_created": 0, "pst_messages_updated": 0, "pst_messages_deleted": 0, "pst_sources_missing": 0, "pst_documents_missing": 0, "mbox_sources_skipped": 0, "mbox_messages_created": 0, "mbox_messages_updated": 0, "mbox_messages_deleted": 0, "mbox_sources_missing": 0, "mbox_documents_missing": 0, "gmail_exports_detected": gmail_export_count, "gmail_linked_documents_created": 0, "gmail_linked_documents_updated": 0, "gmail_documents_scanned": 0, "slack_exports_detected": slack_export_count, "slack_day_documents_scanned": 0, "slack_documents_created": 0, "slack_documents_updated": 0, "slack_documents_missing": 0, "slack_conversations": 0, "email_conversations": 0, "email_documents_reassigned": 0, "email_child_documents_updated": 0, "pst_chat_conversations": 0, "pst_chat_documents_reassigned": 0, "pst_chat_child_documents_updated": 0, "production_documents_created": 0, "production_documents_updated": 0, "production_documents_unchanged": 0, "production_documents_retired": 0, "production_families_reconstructed": 0, "production_docs_missing_linked_text": 0, "production_docs_missing_linked_images": 0, "production_docs_missing_linked_natives": 0, } INGEST_V2_PIPELINE_SCHEMA_VERSION = 1 INGEST_V2_ACTIVE_STATUSES = {"planning", "preparing", "committing", "finalizing"} INGEST_V2_TERMINAL_STATUSES = {"completed", "canceled", "failed"} INGEST_V2_WORK_ITEM_STATUSES = ( "pending", "leased", "prepared", "committing", "committed", "failed", "deferred_timeout", "cancelled", ) INGEST_V2_WORK_ITEM_LEASE_SECONDS = 45 INGEST_V2_PREPARE_BATCH_SIZE = DEFAULT_WORKER_BATCH_SIZE INGEST_V2_PREPARE_MIN_START_SECONDS = 1.0 INGEST_V2_COMMIT_MIN_START_SECONDS = 1.0 INGEST_V2_RUN_STEP_MIN_REMAINING_SECONDS = 1.25 INGEST_V2_RUN_STEP_MAX_INNER_STEPS = 100 INGEST_V2_MAX_SINGLE_STEP_HASH_BYTES = 2 * 1024 * 1024 * 1024 INGEST_V2_BYTES_B64_KEY = "__retriever_bytes_b64__" INGEST_V2_PLAN_CURSOR_SAVE_INTERVAL = 25 INGEST_V2_MBOX_PLAN_BATCH_SIZE = 50 INGEST_V2_MBOX_PLAN_IN_MEMORY_MAX_BYTES = 4 * 1024 * 1024 INGEST_V2_MBOX_INLINE_PAYLOAD_MAX_BYTES = 256 * 1024 INGEST_V2_PREPARED_COMMIT_BATCH_TARGET = max(25, INGEST_V2_PREPARE_BATCH_SIZE * 5) INGEST_V2_PST_PLAN_PAYLOAD_SPILL_DEFAULT_BYTES = 64 * 1024 INGEST_V2_PREPARED_PAYLOAD_SPILL_DEFAULT_BYTES = 24 * 1024 INGEST_V2_MBOX_COMMIT_BATCH_MAX_ITEMS = 16 INGEST_V2_MBOX_COMMIT_BATCH_MAX_PAYLOAD_BYTES = 2 * 1024 * 1024 INGEST_V2_PST_COMMIT_BATCH_MAX_ITEMS = 8 INGEST_V2_PST_COMMIT_BATCH_MAX_PAYLOAD_BYTES = 1 * 1024 * 1024 INGEST_V2_SLACK_COMMIT_BATCH_MAX_ITEMS = 16 INGEST_V2_SLACK_COMMIT_BATCH_MAX_PAYLOAD_BYTES = 2 * 1024 * 1024 INGEST_V2_PRODUCTION_PREVIEW_BATCH_SIZE = 12 INGEST_V2_PRODUCTION_PREVIEW_IMAGE_MAX_DIMENSION = 1400 # Cowork's preview iframe does not reliably resolve sibling local image assets, # so the resumable production path defaults to embedded page images. Flip this # to False to keep HTML references to batch-generated page PNGs. INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES = True INGEST_PIPELINE_V2 = "v2" def ingest_v2_elapsed_ms(started: float) -> float: return (time.perf_counter() - started) * 1000.0 def ingest_v2_percentile_ms(values: list[float], percentile: float) -> float | None: if not values: return None sorted_values = sorted(float(value) for value in values) index = int(((percentile / 100.0) * len(sorted_values)) + 0.999999) - 1 index = max(0, min(len(sorted_values) - 1, index)) return round(sorted_values[index], 3) def ingest_v2_timing_summary(values: list[float]) -> dict[str, object]: if not values: return { "count": 0, "total_ms": 0.0, "avg_ms": None, "p95_ms": None, "max_ms": None, } total = sum(float(value) for value in values) return { "count": len(values), "total_ms": round(total, 3), "avg_ms": round(total / len(values), 3), "p95_ms": ingest_v2_percentile_ms(values, 95.0), "max_ms": round(max(float(value) for value in values), 3), } def ingest_v2_record_worker_event( connection: sqlite3.Connection, *, run_id: str, worker_id: str | None, event_type: str, phase: str, work_item_id: int | None = None, duration_ms: float | None = None, details: dict[str, object] | None = None, created_at: str | None = None, ) -> None: connection.execute( """ INSERT INTO ingest_worker_events ( run_id, worker_id, event_type, work_item_id, phase, duration_ms, details_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, worker_id, event_type, work_item_id, phase, round(float(duration_ms), 3) if duration_ms is not None else None, compact_json_text(ingest_v2_json_safe_value(details or {})), created_at or utc_now(), ), ) def ingest_v2_prepared_payload_spill_threshold_bytes() -> int: raw_value = os.environ.get("RETRIEVER_INGEST_V2_PREPARED_PAYLOAD_SPILL_BYTES") if raw_value is None: return INGEST_V2_PREPARED_PAYLOAD_SPILL_DEFAULT_BYTES raw_value = raw_value.strip() if not raw_value: return INGEST_V2_PREPARED_PAYLOAD_SPILL_DEFAULT_BYTES try: return max(1, int(raw_value)) except ValueError: return INGEST_V2_PREPARED_PAYLOAD_SPILL_DEFAULT_BYTES def ingest_v2_pst_plan_payload_spill_threshold_bytes() -> int: raw_value = os.environ.get("RETRIEVER_INGEST_V2_PST_PLAN_PAYLOAD_SPILL_BYTES") if raw_value is None: return INGEST_V2_PST_PLAN_PAYLOAD_SPILL_DEFAULT_BYTES raw_value = raw_value.strip() if not raw_value: return INGEST_V2_PST_PLAN_PAYLOAD_SPILL_DEFAULT_BYTES try: return max(1, int(raw_value)) except ValueError: return INGEST_V2_PST_PLAN_PAYLOAD_SPILL_DEFAULT_BYTES def ingest_v2_prepared_payload_spill_rel_path(*, run_id: str, work_item_id: int) -> Path: return ( Path("tmp") / "ingest" / "prepared-v2" / sanitize_storage_filename(run_id) / f"{int(work_item_id):08d}.prepared" ) def ingest_v2_prepared_payload_spill_path( paths: dict[str, Path], *, run_id: str, work_item_id: int, ) -> tuple[Path, Path]: rel_path = ingest_v2_prepared_payload_spill_rel_path(run_id=run_id, work_item_id=work_item_id) return rel_path, paths["state_dir"] / rel_path def ingest_v2_pst_plan_payload_spill_rel_path(*, run_id: str, source_key: str) -> Path: return ( Path("payload-spills") / "planned-v2" / sanitize_storage_filename(run_id) / f"{sha256_text(source_key)[:24]}.planned" ) def ingest_v2_pst_plan_payload_spill_path( paths: dict[str, Path], *, run_id: str, source_key: str, ) -> tuple[Path, Path]: rel_path = ingest_v2_pst_plan_payload_spill_rel_path(run_id=run_id, source_key=source_key) return rel_path, paths["state_dir"] / rel_path def ingest_v2_should_spill_prepared_payload( *, payload_kind: str, serialized_payload_bytes: int, ) -> bool: return ( payload_kind == "mbox_message" and serialized_payload_bytes >= ingest_v2_prepared_payload_spill_threshold_bytes() ) def ingest_v2_should_spill_planned_payload( *, unit_type: str, serialized_payload_bytes: int, ) -> bool: return ( unit_type == "pst_message" and serialized_payload_bytes >= ingest_v2_pst_plan_payload_spill_threshold_bytes() ) def ingest_v2_load_spilled_payload(spill_path: Path) -> dict[str, object]: if not spill_path.exists(): raise RetrieverError(f"Spilled ingest payload is missing: {spill_path}") payload = pickle.loads(spill_path.read_bytes()) if isinstance(payload, dict): return payload raise RetrieverError(f"Spilled ingest payload at {spill_path} is malformed.") def ingest_v2_load_spilled_prepared_payload(spill_path: Path) -> dict[str, object]: return ingest_v2_load_spilled_payload(spill_path) def ingest_v2_cleanup_prepared_payload_spills(paths: dict[str, Path], *, run_id: str) -> bool: spill_dir = paths["state_dir"] / ingest_v2_prepared_payload_spill_rel_path(run_id=run_id, work_item_id=0).parent return remove_directory_tree(spill_dir) def ingest_v2_best_effort_cleanup_prepared_payload_spills(paths: dict[str, Path], *, run_id: str) -> None: try: ingest_v2_cleanup_prepared_payload_spills(paths, run_id=run_id) except OSError: return def ingest_v2_cleanup_pst_plan_payload_spills(paths: dict[str, Path], *, run_id: str) -> bool: spill_dir = paths["state_dir"] / ingest_v2_pst_plan_payload_spill_rel_path( run_id=run_id, source_key="cleanup", ).parent return remove_directory_tree(spill_dir) def ingest_v2_best_effort_cleanup_pst_plan_payload_spills(paths: dict[str, Path], *, run_id: str) -> None: try: ingest_v2_cleanup_pst_plan_payload_spills(paths, run_id=run_id) except OSError: return def ingest_v2_attachment_stats(raw_attachments: object) -> tuple[int, int]: attachments = [attachment for attachment in list(raw_attachments or [])] attachment_count = 0 attachment_bytes = 0 for attachment in attachments: attachment_count += 1 payload = attachment.get("payload") if isinstance(attachment, dict) else None if isinstance(payload, bytes): attachment_bytes += len(payload) continue restored_payload = ingest_v2_json_restore_value(payload) if isinstance(restored_payload, bytes): attachment_bytes += len(restored_payload) return attachment_count, attachment_bytes def new_ingest_v2_run_id(now: datetime | None = None) -> str: timestamp = (now or datetime.now(timezone.utc)).strftime("%Y%m%dT%H%M%SZ") return f"{timestamp}-{secrets.token_hex(4)}" def ingest_v2_scope_payload( root: Path, *, recursive: bool, raw_file_types: str | None, raw_paths: list[str] | None, skip_unchanged_loose_files: bool = True, ) -> dict[str, object]: allowed_types = parse_file_types(raw_file_types) scan_scope = build_ingest_scan_scope(root, raw_paths) return { "recursive": bool(recursive), "file_types": sorted(allowed_types) if allowed_types is not None else None, "scan_paths": list(scan_scope.get("display_paths") or []), "skip_unchanged_loose_files": bool(skip_unchanged_loose_files), } def ingest_v2_scan_scope_from_run(root: Path, row: sqlite3.Row) -> dict[str, object]: scope = decode_json_text(row["scope_json"], default={}) or {} raw_paths = list(scope.get("scan_paths") or []) if isinstance(scope, dict) else [] return build_ingest_scan_scope(root, [str(path) for path in raw_paths]) def ingest_v2_cursor_rel_path(root: Path, path: Path) -> str: resolved_root = root.resolve() resolved_path = path.resolve() if resolved_path == resolved_root: return "" return resolved_path.relative_to(resolved_root).as_posix() def ingest_v2_cursor_path(root: Path, rel_path: object) -> Path: normalized = normalize_whitespace(str(rel_path or "")).replace("\\", "/").strip("/") return root if not normalized else root / normalized def ingest_v2_sorted_pending_paths(paths: list[str]) -> list[str]: return sorted(dict.fromkeys(str(path).replace("\\", "/").strip("/") for path in paths)) INGEST_V2_IGNORED_WORKSPACE_PARTS = {".retriever", ".claude"} def ingest_v2_path_is_workspace_internal(root: Path, path: Path) -> bool: try: rel_parts = path.resolve().relative_to(root.resolve()).parts except ValueError: return False return any(part in INGEST_V2_IGNORED_WORKSPACE_PARTS for part in rel_parts) def ingest_v2_initial_pending_paths(root: Path, scan_scope: dict[str, object]) -> list[str]: pending: list[str] = [] for scan_path in list(scan_scope.get("paths") or [root]): pending.append(ingest_v2_cursor_rel_path(root, Path(scan_path))) return ingest_v2_sorted_pending_paths(pending) def ingest_v2_add_excluded_path(root: Path, path: Path, *, exact: set[str], prefixes: set[str]) -> None: if not path_is_at_or_under(path, root): return rel_path = ingest_v2_cursor_rel_path(root, path) if not rel_path: return exact.add(rel_path) if path.is_dir(): prefixes.add(rel_path.rstrip("/") + "/") def ingest_v2_gmail_drive_record_payload(root: Path, record: dict[str, object]) -> dict[str, object]: payload = dict(record) file_path = payload.pop("file_path", None) if isinstance(file_path, Path) and path_is_at_or_under(file_path, root): payload["file_rel_path"] = relative_document_path(root, file_path) return ingest_v2_json_safe_value(payload) def ingest_v2_gmail_drive_record_from_payload(root: Path, record: dict[str, object]) -> dict[str, object]: restored = dict(record) file_rel_path = normalize_whitespace(str(restored.pop("file_rel_path", "") or "")) if file_rel_path: restored["file_path"] = ingest_v2_cursor_path(root, file_rel_path) return restored def ingest_v2_gmail_mbox_source_payloads( root: Path, descriptors: list[dict[str, object]], ) -> list[dict[str, object]]: source_payloads: list[dict[str, object]] = [] for descriptor in descriptors: root_rel_path = ingest_v2_cursor_rel_path(root, Path(descriptor["root"])) linked_drive_attachment_records_by_message_id = { str(message_id): [ ingest_v2_gmail_drive_record_payload(root, dict(record)) for record in list(records) ] for message_id, records in dict(descriptor.get("linked_drive_attachment_records_by_message_id") or {}).items() } linked_drive_rel_paths = sorted( { str(record.get("file_rel_path")) for records in linked_drive_attachment_records_by_message_id.values() for record in records if normalize_whitespace(str(record.get("file_rel_path") or "")) } ) for mbox_path in sorted([Path(path) for path in list(descriptor.get("mbox_paths") or [])], key=lambda path: path.as_posix()): source_payloads.append( { "source_plan_kind": "gmail", "gmail_export_root_rel_path": root_rel_path, "source_rel_path": relative_document_path(root, mbox_path), "message_sidecar_hash": normalize_whitespace( str(descriptor.get("message_sidecar_hash") or "") ) or None, "email_metadata_by_message_id": ingest_v2_json_safe_value( dict(descriptor.get("email_metadata_by_message_id") or {}) ), "linked_drive_records_by_message_id": ingest_v2_json_safe_value( dict(descriptor.get("linked_drive_records_by_message_id") or {}) ), "linked_drive_attachment_records_by_message_id": linked_drive_attachment_records_by_message_id, "linked_drive_rel_paths": linked_drive_rel_paths, } ) return source_payloads def ingest_v2_gmail_mbox_source_scan_hash(path: Path, source_payload: dict[str, object]) -> str: return sha256_json_value( { "mbox_hash": sha256_file(path), "message_sidecar_hash": normalize_whitespace( str(source_payload.get("message_sidecar_hash") or "") ) or None, "source_rel_path": str(source_payload["source_rel_path"]), } ) def ingest_v2_pst_export_source_payloads_by_rel_path( root: Path, descriptors: list[dict[str, object]], ) -> dict[str, dict[str, object]]: payloads: dict[str, dict[str, object]] = {} for descriptor in descriptors: message_metadata_by_pst_path = dict(descriptor.get("message_metadata_by_pst_path") or {}) message_match_records_by_pst_path = dict(descriptor.get("message_match_records_by_pst_path") or {}) message_sidecar_hash = normalize_whitespace(str(descriptor.get("message_sidecar_hash") or "")) or None for pst_path in sorted([Path(path) for path in list(descriptor.get("pst_paths") or [])], key=lambda path: path.as_posix()): resolved_key = pst_path.resolve().as_posix() rel_path = relative_document_path(root, pst_path) payloads[rel_path] = { "source_plan_kind": "pst_export", "source_rel_path": rel_path, "message_sidecar_hash": message_sidecar_hash, "message_metadata_by_source_item": ingest_v2_json_safe_value( dict(message_metadata_by_pst_path.get(resolved_key) or {}) ), "message_match_records": ingest_v2_json_safe_value( list(message_match_records_by_pst_path.get(resolved_key) or []) ), } return payloads def ingest_v2_pst_source_scan_hash(path: Path, source_payload: dict[str, object] | None = None) -> str: payload = dict(source_payload or {}) message_sidecar_hash = normalize_whitespace(str(payload.get("message_sidecar_hash") or "")) or None if message_sidecar_hash: return sha256_json_value( { "pst_hash": sha256_file(path), "message_sidecar_hash": message_sidecar_hash, "sidecar_match_version": "pst-export-sidecar-v2", "source_rel_path": str(payload.get("source_rel_path") or ""), } ) return sha256_text(f"pst-ingest-v5:{sha256_file(path) or ''}") def ingest_v2_slack_export_descriptor_payload(root: Path, descriptor: dict[str, object]) -> dict[str, object]: export_root = Path(descriptor["root"]).resolve() return { "rel_root": relative_document_path(root, export_root), "day_rel_paths": [ relative_document_path(root, Path(day_file)) for day_file in list(descriptor.get("day_files") or []) ], } def ingest_v2_planning_exclusions( root: Path, recursive: bool, allowed_types: set[str] | None, connection: sqlite3.Connection, scan_scope: dict[str, object], ) -> dict[str, object]: exact: set[str] = set() prefixes: set[str] = set() production_signatures = find_production_root_signatures(root, recursive, connection, scan_scope=scan_scope) for signature in production_signatures: ingest_v2_add_excluded_path(root, Path(signature["root"]), exact=exact, prefixes=prefixes) slack_export_descriptors = find_scoped_source_roots( find_slack_export_roots, root, recursive, scan_scope, allowed_types, ) for descriptor in slack_export_descriptors: ingest_v2_add_excluded_path(root, Path(descriptor["root"]), exact=exact, prefixes=prefixes) slack_export_payloads = [ ingest_v2_slack_export_descriptor_payload(root, descriptor) for descriptor in slack_export_descriptors ] gmail_export_descriptors = find_scoped_source_roots( find_gmail_export_roots, root, recursive, scan_scope, allowed_types, ) for descriptor in gmail_export_descriptors: ingest_v2_add_excluded_path(root, Path(descriptor["root"]), exact=exact, prefixes=prefixes) for owned_path in list(descriptor.get("owned_paths") or []): ingest_v2_add_excluded_path(root, Path(owned_path), exact=exact, prefixes=prefixes) gmail_mbox_source_payloads = ( ingest_v2_gmail_mbox_source_payloads(root, gmail_export_descriptors) if allowed_types is None or MBOX_SOURCE_KIND in allowed_types else [] ) pst_export_descriptors = ( find_scoped_source_roots(find_pst_export_roots, root, recursive, scan_scope) if allowed_types is None or PST_SOURCE_KIND in allowed_types else [] ) pst_source_payloads_by_rel_path = ingest_v2_pst_export_source_payloads_by_rel_path(root, pst_export_descriptors) for descriptor in pst_export_descriptors: for owned_path in list(descriptor.get("owned_paths") or []): ingest_v2_add_excluded_path(root, Path(owned_path), exact=exact, prefixes=prefixes) return { "exact_rel_paths": sorted(exact), "dir_prefixes": sorted(prefixes), "slack_export_payloads": slack_export_payloads, "gmail_mbox_source_payloads": gmail_mbox_source_payloads, "pst_source_payloads_by_rel_path": pst_source_payloads_by_rel_path, "counts": { "production_roots": len(production_signatures), "slack_export_roots": len(slack_export_descriptors), "gmail_export_roots": len(gmail_export_descriptors), "pst_export_roots": len(pst_export_descriptors), }, } def ingest_v2_rel_path_excluded(cursor: dict[str, object], rel_path: str) -> bool: if not rel_path: return False exact = set(cursor.get("excluded_exact_rel_paths") or []) prefixes = list(cursor.get("excluded_dir_prefixes") or []) return rel_path in exact or any(rel_path.startswith(str(prefix)) for prefix in prefixes) def ingest_v2_save_phase_cursor( connection: sqlite3.Connection, *, run_id: str, phase: str, cursor_key: str, cursor: dict[str, object], status: str, ) -> None: now = utc_now() connection.execute( """ INSERT INTO ingest_phase_cursors ( run_id, phase, cursor_key, cursor_json, status, updated_at ) VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(run_id, phase, cursor_key) DO UPDATE SET cursor_json = excluded.cursor_json, status = excluded.status, updated_at = excluded.updated_at """, (run_id, phase, cursor_key, compact_json_text(cursor), status, now), ) def ingest_v2_save_planning_cursor_heartbeat( connection: sqlite3.Connection, *, run_id: str, cursor: dict[str, object], status: str = "pending", ) -> float: started = time.perf_counter() if not connection.in_transaction: connection.execute("BEGIN") try: ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="planning", cursor_key="loose_file_scan", cursor=cursor, status=status, ) connection.execute( """ UPDATE ingest_runs SET last_heartbeat_at = ? WHERE run_id = ? """, (utc_now(), run_id), ) connection.commit() except Exception: connection.rollback() raise return ingest_v2_elapsed_ms(started) def ingest_v2_load_or_create_loose_file_plan_cursor( connection: sqlite3.Connection, root: Path, row: sqlite3.Row, ) -> dict[str, object]: run_id = str(row["run_id"]) cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'planning' AND cursor_key = 'loose_file_scan' """, (run_id,), ).fetchone() if cursor_row is not None: cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if isinstance(cursor, dict): cursor.setdefault("pending_production_rel_roots", []) cursor.setdefault("production_roots_by_rel_root", {}) cursor.setdefault("planned_production_roots", []) cursor.setdefault("skipped_production_roots", []) cursor.setdefault("planned_production_rows", 0) cursor.setdefault("production_failures", []) cursor.setdefault("production_docs_missing_linked_text", 0) cursor.setdefault("production_docs_missing_linked_images", 0) cursor.setdefault("production_docs_missing_linked_natives", 0) cursor.setdefault("pending_slack_export_roots", []) cursor.setdefault("slack_export_roots_by_rel_root", {}) cursor.setdefault("planned_slack_export_roots", []) cursor.setdefault("planned_slack_conversations", 0) cursor.setdefault("planned_slack_day_documents", 0) cursor.setdefault("planned_slack_day_files_scanned", 0) cursor.setdefault("slack_failures", []) cursor.setdefault("current_mbox_source", None) cursor.setdefault("pending_gmail_mbox_sources", []) cursor.setdefault("planned_mbox_sources", []) cursor.setdefault("planned_mbox_messages", 0) cursor.setdefault("planned_gmail_mbox_sources", 0) cursor.setdefault("skipped_mbox_sources", 0) cursor.setdefault("scanned_mbox_source_rel_paths", []) cursor.setdefault("mbox_failures", []) cursor.setdefault("current_pst_source", None) cursor.setdefault("pst_source_payloads_by_rel_path", {}) cursor.setdefault("planned_pst_sources", []) cursor.setdefault("planned_pst_messages", 0) cursor.setdefault("skipped_pst_sources", 0) cursor.setdefault("scanned_pst_source_rel_paths", []) cursor.setdefault("pst_failures", []) cursor.setdefault("skipped_unchanged_loose_files", 0) cursor.setdefault("skipped_unchanged_loose_file_rel_paths", []) return cursor recursive = bool(row["recursive"]) allowed_types = parse_file_types(row["raw_file_types"]) scan_scope = ingest_v2_scan_scope_from_run(root, row) exclusions = ingest_v2_planning_exclusions(root, recursive, allowed_types, connection, scan_scope) production_signatures = find_production_root_signatures(root, recursive, connection, scan_scope=scan_scope) production_payloads = [ ingest_v2_production_signature_payload(root, signature) for signature in production_signatures ] next_order_row = connection.execute( """ SELECT COALESCE(MAX(commit_order), 0) + 1 AS next_order FROM ingest_work_items WHERE run_id = ? """, (run_id,), ).fetchone() cursor = { "schema_version": 1, "pending_paths": ingest_v2_initial_pending_paths(root, scan_scope), "next_commit_order": int(next_order_row["next_order"] or 1), "scanned_paths": 0, "planned_loose_files": 0, "skipped_unchanged_loose_files": 0, "skipped_unchanged_loose_file_rel_paths": [], "skipped_container_files": 0, "skipped_extensionless_files": 0, "skipped_filtered_files": 0, "skipped_excluded_paths": 0, "skipped_missing_paths": 0, "listed_directories": 0, "pending_production_rel_roots": ( [str(payload["rel_root"]) for payload in production_payloads] if allowed_types is None else [] ), "production_roots_by_rel_root": { str(payload["rel_root"]): payload for payload in production_payloads }, "planned_production_roots": [], "skipped_production_roots": ( [str(payload["rel_root"]) for payload in production_payloads] if allowed_types is not None else [] ), "planned_production_rows": 0, "production_failures": [], "production_docs_missing_linked_text": 0, "production_docs_missing_linked_images": 0, "production_docs_missing_linked_natives": 0, "pending_slack_export_roots": [str(payload["rel_root"]) for payload in list(exclusions.get("slack_export_payloads") or [])], "slack_export_roots_by_rel_root": { str(payload["rel_root"]): payload for payload in list(exclusions.get("slack_export_payloads") or []) }, "planned_slack_export_roots": [], "planned_slack_conversations": 0, "planned_slack_day_documents": 0, "planned_slack_day_files_scanned": 0, "slack_failures": [], "current_mbox_source": None, "pending_gmail_mbox_sources": list(exclusions.get("gmail_mbox_source_payloads") or []), "planned_mbox_sources": [], "planned_mbox_messages": 0, "planned_gmail_mbox_sources": 0, "skipped_mbox_sources": 0, "scanned_mbox_source_rel_paths": [], "mbox_failures": [], "current_pst_source": None, "pst_source_payloads_by_rel_path": dict(exclusions.get("pst_source_payloads_by_rel_path") or {}), "planned_pst_sources": [], "planned_pst_messages": 0, "skipped_pst_sources": 0, "scanned_pst_source_rel_paths": [], "pst_failures": [], "excluded_exact_rel_paths": list(exclusions["exact_rel_paths"]), "excluded_dir_prefixes": list(exclusions["dir_prefixes"]), "special_source_counts": exclusions["counts"], } connection.execute("BEGIN") try: ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="planning", cursor_key="loose_file_scan", cursor=cursor, status="pending", ) connection.commit() except Exception: connection.rollback() raise return cursor def ingest_v2_plan_loose_file_item( connection: sqlite3.Connection, *, run_id: str, rel_path: str, file_type: str, file_size: int | None, file_mtime_ns: int | None, commit_order: int, ) -> bool: now = utc_now() payload = { "rel_path": rel_path, "file_type": file_type, "source_file_size": file_size, "source_file_mtime_ns": file_mtime_ns, "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "loose_file", FILESYSTEM_SOURCE_KIND, rel_path, rel_path, int(commit_order), compact_json_text(payload), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_loose_file_matches_existing_snapshot( connection: sqlite3.Connection, path: Path, *, rel_path: str, file_size: int | None, ) -> bool: if file_size is None or file_size > INGEST_V2_MAX_SINGLE_STEP_HASH_BYTES: return False occurrence_row = connection.execute( """ SELECT * FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND rel_path = ? AND lifecycle_status = ? ORDER BY id ASC LIMIT 1 """, (FILESYSTEM_SOURCE_KIND, rel_path, ACTIVE_OCCURRENCE_STATUS), ).fetchone() if occurrence_row is None: return False if occurrence_row["file_size"] is None or int(occurrence_row["file_size"]) != int(file_size): return False stored_hash = normalize_whitespace(str(occurrence_row["file_hash"] or "")) if not stored_hash: return False document_row = connection.execute( """ SELECT * FROM documents WHERE id = ? AND lifecycle_status != 'deleted' """, (int(occurrence_row["document_id"]),), ).fetchone() if document_row is None or document_row["dataset_id"] is None or occurrence_row["dataset_source_id"] is None: return False if not document_row_has_seeded_text_revisions(document_row): return False if not document_row_has_email_threading(connection, document_row): return False try: current_hash = sha256_file(path) except OSError: return False return current_hash == stored_hash def ingest_v2_plan_slack_conversation_item( connection: sqlite3.Connection, *, run_id: str, rel_root: str, conversation_plan: dict[str, object], commit_order: int, ) -> bool: now = utc_now() conversation_key = str(conversation_plan["conversation_key"]) rel_paths = [str(rel_path) for rel_path in list(conversation_plan.get("rel_paths") or [])] payload = { **ingest_v2_json_safe_value(conversation_plan), "source_locator": rel_root, "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, affected_conversation_keys_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "slack_conversation", SLACK_EXPORT_SOURCE_KIND, f"{rel_root}:{conversation_key}", rel_paths[0] if rel_paths else rel_root, int(commit_order), compact_json_text(payload), compact_json_text([f"{rel_root}:{conversation_key}"]), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_plan_slack_document_item( connection: sqlite3.Connection, *, run_id: str, rel_root: str, conversation_plan: dict[str, object], document_plan: dict[str, object], commit_order: int, conversation_lead_document: bool, ) -> bool: now = utc_now() conversation_key = str(conversation_plan["conversation_key"]) plan = dict(document_plan.get("plan") or {}) rel_path = str(plan.get("rel_path") or "") payload = { **ingest_v2_json_safe_value(document_plan), "document_kind": str(document_plan.get("kind") or ""), "source_locator": rel_root, "conversation_key": conversation_key, "conversation_type": str(conversation_plan["conversation_type"]), "display_name": str(conversation_plan["display_name"]), "conversation_lead_document": bool(conversation_lead_document), "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, affected_conversation_keys_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "slack_document", SLACK_EXPORT_SOURCE_KIND, f"{rel_root}:{rel_path}", rel_path, int(commit_order), compact_json_text(payload), compact_json_text([f"{rel_root}:{conversation_key}"]), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_plan_slack_export_root( connection: sqlite3.Connection, root: Path, *, run_id: str, payload: dict[str, object], next_commit_order: int, ) -> dict[str, object]: rel_root = str(payload["rel_root"]) export_root = ingest_v2_cursor_path(root, rel_root) conversation_directory = load_slack_export_conversation_directory(export_root) user_directory = load_slack_user_directory(export_root) day_files = [ ingest_v2_cursor_path(root, str(day_rel_path)) for day_rel_path in list(payload.get("day_rel_paths") or []) ] if not day_files: day_files = iter_slack_export_day_files(export_root) conversation_plans = plan_slack_export_conversations( root, export_root, conversation_directory=conversation_directory, user_directory=user_directory, day_files=day_files, ) commit_order = int(next_commit_order) planned_conversations = 0 planned_day_documents = 0 planned_day_files_scanned = len(day_files) rel_paths: list[str] = [] for conversation_plan in conversation_plans: conversation_inserted = False ordered_documents = [ *list(conversation_plan.get("day_documents") or []), *list(conversation_plan.get("thread_documents") or []), ] for document_index, document_plan in enumerate(ordered_documents): inserted = ingest_v2_plan_slack_document_item( connection, run_id=run_id, rel_root=rel_root, conversation_plan=conversation_plan, document_plan=document_plan, commit_order=commit_order, conversation_lead_document=document_index == 0, ) commit_order += 1 if inserted: conversation_inserted = True if str(document_plan.get("kind") or "") == "day": planned_day_documents += 1 if conversation_inserted: planned_conversations += 1 rel_paths.extend(str(rel_path) for rel_path in list(conversation_plan.get("rel_paths") or [])) return { "rel_root": rel_root, "next_commit_order": commit_order, "planned_conversations": planned_conversations, "planned_day_documents": planned_day_documents, "planned_day_files_scanned": planned_day_files_scanned, "seen_rel_paths": sorted(set(rel_paths)), } def ingest_v2_mbox_source_scan_hash(path: Path) -> str: return sha256_text(f"mbox-ingest-v1:{sha256_file(path) or ''}") def ingest_v2_mbox_blank_separator_line(line: bytes) -> bool: return line in (b"\n", b"\r\n") def ingest_v2_split_small_mbox_messages(data: bytes) -> list[bytes]: if not data: return [] messages: list[bytes] = [] current_lines: list[bytes] | None = None for line in data.splitlines(keepends=True): if line.startswith(b"From "): if current_lines is not None: payload_lines = ( current_lines[:-1] if current_lines and ingest_v2_mbox_blank_separator_line(current_lines[-1]) else current_lines ) messages.append(b"".join(payload_lines)) current_lines = [] continue if current_lines is None: if line.strip(): raise RetrieverError("MBOX planner could not find a leading 'From ' separator line.") continue current_lines.append(line) if current_lines is None: if data.strip(): raise RetrieverError("MBOX planner found content but no message separators.") return [] payload_lines = ( current_lines[:-1] if current_lines and ingest_v2_mbox_blank_separator_line(current_lines[-1]) else current_lines ) messages.append(b"".join(payload_lines)) return messages def ingest_v2_iter_mbox_plan_messages(path: Path) -> Iterator[dict[str, object]]: file_size = file_size_bytes(path) if file_size is not None and 0 <= int(file_size) <= INGEST_V2_MBOX_PLAN_IN_MEMORY_MAX_BYTES: try: parser = BytesParser(policy=policy.default) for message_index, raw_payload in enumerate(ingest_v2_split_small_mbox_messages(path.read_bytes())): parsed_message = parser.parsebytes(raw_payload) payload_bytes = parsed_message.as_bytes(policy=policy.default, unixfrom=False) yield { "message_index": int(message_index), "message_key": int(message_index), "payload_bytes": payload_bytes, "explicit_source_item_id": normalize_whitespace( str(parsed_message.get("Message-ID") or parsed_message.get("Message-Id") or "") ) or None, "parser_kind": "in_memory", } return except Exception: pass archive = mailbox.mbox(str(path), factory=mailbox.mboxMessage, create=False) try: for message_index, (message_key, raw_message) in enumerate(archive.iteritems()): payload_bytes = raw_message.as_bytes(policy=policy.default, unixfrom=False) yield { "message_index": int(message_index), "message_key": message_key, "payload_bytes": payload_bytes, "explicit_source_item_id": normalize_whitespace( str(raw_message.get("Message-ID") or raw_message.get("Message-Id") or "") ) or None, "parser_kind": "mailbox", } finally: try: archive.close() except Exception: pass def ingest_v2_plan_mbox_message_item( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, source_plan_kind: str, message_index: int, message_key: object, source_item_id: str, payload_hash: str, source_file_size: int | None, source_file_mtime: str | None, source_file_hash: str, scan_started_at: str, commit_order: int, message_metadata: dict[str, object] | None = None, linked_drive_records: list[dict[str, object]] | None = None, linked_drive_attachment_records: list[dict[str, object]] | None = None, payload_bytes: bytes | None = None, ) -> bool: now = utc_now() rel_path = mbox_message_rel_path(source_rel_path, source_item_id) payload = { "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "message_index": int(message_index), "message_key": message_key, "source_item_id": source_item_id, "payload_hash": payload_hash, "source_file_size": source_file_size, "source_file_mtime": source_file_mtime, "source_file_hash": source_file_hash, "scan_started_at": scan_started_at, "planned_at": now, } if payload_bytes is not None: payload["payload_bytes"] = payload_bytes if message_metadata: payload["message_metadata"] = dict(message_metadata) if linked_drive_records: payload["linked_drive_records"] = list(linked_drive_records) if linked_drive_attachment_records: payload["linked_drive_attachment_records"] = list(linked_drive_attachment_records) cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, parent_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "mbox_message", MBOX_SOURCE_KIND, f"{source_rel_path}:{source_item_id}", rel_path, int(commit_order), int(message_index), compact_json_text(ingest_v2_json_safe_value(payload)), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_plan_mbox_source_finalizer_item( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, source_file_size: int | None, source_file_mtime: str | None, source_file_hash: str, scan_started_at: str, message_count: int, skip_source: bool, commit_order: int, source_plan_kind: str = "mbox", source_action: str | None = None, linked_drive_rel_paths: list[str] | None = None, ) -> bool: now = utc_now() normalized_source_action = normalize_whitespace(str(source_action or "")) or ( "skipped" if skip_source else "updated" ) payload = { "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "source_action": normalized_source_action, "source_file_size": source_file_size, "source_file_mtime": source_file_mtime, "source_file_hash": source_file_hash, "scan_started_at": scan_started_at, "message_count": int(message_count), "skip_source": bool(skip_source), "linked_drive_rel_paths": list(linked_drive_rel_paths or []), "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "mbox_source_finalizer", MBOX_SOURCE_KIND, f"{source_rel_path}:__finalize__", source_rel_path, int(commit_order), compact_json_text(ingest_v2_json_safe_value(payload)), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_existing_mbox_message_count( connection: sqlite3.Connection, *, source_rel_path: str, ) -> int: source_row = get_container_source_row(connection, MBOX_SOURCE_KIND, source_rel_path) if source_row is not None and int(source_row["message_count"] or 0) > 0: return int(source_row["message_count"] or 0) return len( container_root_occurrence_rows_for_source( connection, source_kind=MBOX_SOURCE_KIND, source_rel_path=source_rel_path, ) ) def ingest_v2_existing_pst_message_count( connection: sqlite3.Connection, *, source_rel_path: str, ) -> int: source_row = get_container_source_row(connection, PST_SOURCE_KIND, source_rel_path) if source_row is not None and int(source_row["message_count"] or 0) > 0: return int(source_row["message_count"] or 0) return len( container_root_occurrence_rows_for_source( connection, source_kind=PST_SOURCE_KIND, source_rel_path=source_rel_path, ) ) def ingest_v2_begin_mbox_source_plan( connection: sqlite3.Connection, root: Path, *, run_id: str, rel_path: str, commit_order: int, source_plan_kind: str = "mbox", source_scan_hash: str | None = None, gmail_source_payload: dict[str, object] | None = None, ) -> tuple[dict[str, object] | None, int, bool]: path = ingest_v2_cursor_path(root, rel_path) source_file_size = file_size_bytes(path) source_file_mtime = file_mtime_timestamp(path) source_file_hash = source_scan_hash or ingest_v2_mbox_source_scan_hash(path) gmail_payload = dict(gmail_source_payload or {}) existing_source = get_container_source_row(connection, MBOX_SOURCE_KIND, rel_path) scan_started_at = next_monotonic_utc_timestamp( [ existing_source["last_scan_started_at"] if existing_source is not None else None, existing_source["last_scan_completed_at"] if existing_source is not None else None, ] ) if ( existing_source is not None and container_source_scan_completed(existing_source) and not container_documents_missing_text_revisions( connection, source_kind=MBOX_SOURCE_KIND, source_rel_path=rel_path, ) and not container_email_documents_missing_threading( connection, source_kind=MBOX_SOURCE_KIND, source_rel_path=rel_path, ) and existing_source["file_size"] == source_file_size and existing_source["file_hash"] == source_file_hash and (existing_source["file_mtime"] == source_file_mtime or existing_source["file_hash"]) ): message_count = ingest_v2_existing_mbox_message_count(connection, source_rel_path=rel_path) ingest_v2_plan_mbox_source_finalizer_item( connection, run_id=run_id, source_rel_path=rel_path, source_plan_kind=source_plan_kind, source_file_size=source_file_size, source_file_mtime=source_file_mtime, source_file_hash=source_file_hash, scan_started_at=scan_started_at, message_count=message_count, skip_source=True, commit_order=commit_order, source_action="skipped", linked_drive_rel_paths=list(gmail_payload.get("linked_drive_rel_paths") or []), ) return None, commit_order + 1, True return ( { "source_rel_path": rel_path, "source_plan_kind": source_plan_kind, "source_action": "new" if existing_source is None else "updated", "source_file_size": source_file_size, "source_file_mtime": source_file_mtime, "source_file_hash": source_file_hash, "scan_started_at": scan_started_at, "next_message_index": 0, "planned_message_count": 0, "duplicate_source_item_counts": {}, "next_commit_order": int(commit_order), "email_metadata_by_message_id": dict(gmail_payload.get("email_metadata_by_message_id") or {}), "linked_drive_records_by_message_id": dict(gmail_payload.get("linked_drive_records_by_message_id") or {}), "linked_drive_attachment_records_by_message_id": dict( gmail_payload.get("linked_drive_attachment_records_by_message_id") or {} ), "linked_drive_rel_paths": list(gmail_payload.get("linked_drive_rel_paths") or []), }, commit_order, False, ) def ingest_v2_plan_current_mbox_source( connection: sqlite3.Connection, root: Path, *, run_id: str, current_mbox_source: dict[str, object], deadline: float, ) -> tuple[dict[str, object] | None, int, int, bool, dict[str, object]]: source_rel_path = str(current_mbox_source["source_rel_path"]) source_plan_kind = str(current_mbox_source.get("source_plan_kind") or "mbox") path = ingest_v2_cursor_path(root, source_rel_path) next_message_index = int(current_mbox_source.get("next_message_index") or 0) next_commit_order = int(current_mbox_source.get("next_commit_order") or 1) processed_this_step = 0 plan_started = time.perf_counter() planned_payload_bytes = 0 planned_inline_payload_bytes = 0 parser_kind_counts: dict[str, int] = {} message_index_start: int | None = None message_index_end: int | None = None duplicate_counts = { str(key): int(value) for key, value in dict(current_mbox_source.get("duplicate_source_item_counts") or {}).items() } reached_end = True for plan_message in ingest_v2_iter_mbox_plan_messages(path): message_index = int(plan_message["message_index"]) if message_index < next_message_index: continue if ( processed_this_step >= INGEST_V2_MBOX_PLAN_BATCH_SIZE or ingest_v2_deadline_remaining_seconds(deadline) < 1.0 ): reached_end = False break payload_bytes = bytes(plan_message["payload_bytes"]) planned_payload_bytes += len(payload_bytes) parser_kind = str(plan_message.get("parser_kind") or "mailbox") parser_kind_counts[parser_kind] = int(parser_kind_counts.get(parser_kind) or 0) + 1 if message_index_start is None: message_index_start = message_index message_index_end = message_index payload_hash = sha256_bytes(payload_bytes) explicit_source_item_id = normalize_whitespace(str(plan_message.get("explicit_source_item_id") or "")) or None base_source_item_id = explicit_source_item_id or f"mbox-hash:{payload_hash}" duplicate_counts[base_source_item_id] = int(duplicate_counts.get(base_source_item_id) or 0) + 1 occurrence = int(duplicate_counts[base_source_item_id]) source_item_id = base_source_item_id if occurrence == 1 else f"{base_source_item_id}#{occurrence}" message_lookup_key = gmail_normalized_message_lookup_key(source_item_id) message_metadata = ( dict(dict(current_mbox_source.get("email_metadata_by_message_id") or {}).get(message_lookup_key) or {}) if message_lookup_key is not None else {} ) linked_drive_records = ( list(dict(current_mbox_source.get("linked_drive_records_by_message_id") or {}).get(message_lookup_key) or []) if message_lookup_key is not None else [] ) linked_drive_attachment_records = ( list( dict(current_mbox_source.get("linked_drive_attachment_records_by_message_id") or {}).get( message_lookup_key ) or [] ) if message_lookup_key is not None else [] ) inline_payload_bytes = ( payload_bytes if len(payload_bytes) <= INGEST_V2_MBOX_INLINE_PAYLOAD_MAX_BYTES else None ) if inline_payload_bytes is not None: planned_inline_payload_bytes += len(inline_payload_bytes) ingest_v2_plan_mbox_message_item( connection, run_id=run_id, source_rel_path=source_rel_path, source_plan_kind=source_plan_kind, message_index=message_index, message_key=plan_message.get("message_key"), source_item_id=source_item_id, payload_hash=payload_hash, source_file_size=( int(current_mbox_source["source_file_size"]) if current_mbox_source.get("source_file_size") is not None else None ), source_file_mtime=( str(current_mbox_source["source_file_mtime"]) if current_mbox_source.get("source_file_mtime") is not None else None ), source_file_hash=str(current_mbox_source["source_file_hash"]), scan_started_at=str(current_mbox_source["scan_started_at"]), commit_order=next_commit_order, message_metadata=message_metadata, linked_drive_records=linked_drive_records, linked_drive_attachment_records=linked_drive_attachment_records, payload_bytes=inline_payload_bytes, ) processed_this_step += 1 next_commit_order += 1 next_message_index = message_index + 1 current_mbox_source["next_message_index"] = next_message_index current_mbox_source["next_commit_order"] = next_commit_order current_mbox_source["planned_message_count"] = ( int(current_mbox_source.get("planned_message_count") or 0) + processed_this_step ) current_mbox_source["duplicate_source_item_counts"] = duplicate_counts plan_metrics = { "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "planned_messages": processed_this_step, "payload_bytes": planned_payload_bytes, "inline_payload_bytes": planned_inline_payload_bytes, "parser_kind_counts": parser_kind_counts, "message_index_start": message_index_start, "message_index_end": message_index_end, "duration_ms": ingest_v2_elapsed_ms(plan_started), "source_complete": reached_end, } if not reached_end: return current_mbox_source, next_commit_order, processed_this_step, False, plan_metrics ingest_v2_plan_mbox_source_finalizer_item( connection, run_id=run_id, source_rel_path=source_rel_path, source_plan_kind=source_plan_kind, source_action=str(current_mbox_source.get("source_action") or "updated"), source_file_size=( int(current_mbox_source["source_file_size"]) if current_mbox_source.get("source_file_size") is not None else None ), source_file_mtime=( str(current_mbox_source["source_file_mtime"]) if current_mbox_source.get("source_file_mtime") is not None else None ), source_file_hash=str(current_mbox_source["source_file_hash"]), scan_started_at=str(current_mbox_source["scan_started_at"]), message_count=int(current_mbox_source.get("planned_message_count") or 0), skip_source=False, commit_order=next_commit_order, linked_drive_rel_paths=list(current_mbox_source.get("linked_drive_rel_paths") or []), ) return None, next_commit_order + 1, processed_this_step, True, plan_metrics def ingest_v2_plan_pst_message_item( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, source_plan_kind: str, message_index: int, raw_message: dict[str, object], source_item_id: str, source_file_size: int | None, source_file_mtime: str | None, source_file_hash: str, scan_started_at: str, commit_order: int, message_metadata: dict[str, object] | None = None, message_match_records: list[dict[str, object]] | None = None, ) -> tuple[bool, dict[str, object]]: now = utc_now() rel_path = pst_message_rel_path(source_rel_path, source_item_id) source_key = f"{source_rel_path}:{source_item_id}" payload = { "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "message_index": int(message_index), "source_item_id": source_item_id, "source_file_size": source_file_size, "source_file_mtime": source_file_mtime, "source_file_hash": source_file_hash, "scan_started_at": scan_started_at, "planned_at": now, } active_paths = active_workspace_paths() raw_message_dict = dict(raw_message) raw_message_spill_path: Path | None = None serialized_spill_payload: bytes | None = None payload_bytes = 0 spilled_payload_bytes = 0 try: if active_paths is not None: serialized_spill_payload = serialize_prepared_ingest_item({"raw_message": raw_message_dict}) if ingest_v2_should_spill_planned_payload( unit_type="pst_message", serialized_payload_bytes=len(serialized_spill_payload), ): spill_rel_path_obj, raw_message_spill_path = ingest_v2_pst_plan_payload_spill_path( active_paths, run_id=run_id, source_key=source_key, ) raw_message_spill_path.parent.mkdir(parents=True, exist_ok=True) raw_message_spill_path.write_bytes(serialized_spill_payload) payload["raw_message_spilled"] = True payload["raw_message_spill_rel_path"] = spill_rel_path_obj.as_posix() payload["raw_message_payload_bytes"] = len(serialized_spill_payload) payload_bytes = len(serialized_spill_payload) spilled_payload_bytes = len(serialized_spill_payload) else: payload["raw_message"] = ingest_v2_json_safe_value(raw_message_dict) else: payload["raw_message"] = ingest_v2_json_safe_value(raw_message_dict) if "raw_message" not in payload and not payload.get("raw_message_spilled"): payload["raw_message"] = ingest_v2_json_safe_value(raw_message_dict) if message_metadata: payload["message_metadata_by_source_item"] = {source_item_id: dict(message_metadata)} if message_match_records: payload["message_match_records"] = list(message_match_records) payload_json = compact_json_text(ingest_v2_json_safe_value(payload)) if not payload.get("raw_message_spilled"): payload_bytes = len(payload_json.encode("utf-8")) cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, parent_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "pst_message", PST_SOURCE_KIND, source_key, rel_path, int(commit_order), int(message_index), payload_json, "pending", now, now, ), ) inserted = int(cursor.rowcount or 0) > 0 if not inserted and raw_message_spill_path is not None: remove_file_if_exists(raw_message_spill_path) return inserted, { "payload_bytes": payload_bytes if inserted else 0, "spilled_payload_bytes": spilled_payload_bytes if inserted else 0, "spilled": bool(inserted and raw_message_spill_path is not None), } except Exception: if raw_message_spill_path is not None: remove_file_if_exists(raw_message_spill_path) raise def ingest_v2_plan_pst_source_finalizer_item( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, source_file_size: int | None, source_file_mtime: str | None, source_file_hash: str, scan_started_at: str, message_count: int, skip_source: bool, commit_order: int, source_plan_kind: str = "pst", source_action: str | None = None, ) -> bool: now = utc_now() normalized_source_action = normalize_whitespace(str(source_action or "")) or ( "skipped" if skip_source else "updated" ) payload = { "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "source_action": normalized_source_action, "source_file_size": source_file_size, "source_file_mtime": source_file_mtime, "source_file_hash": source_file_hash, "scan_started_at": scan_started_at, "message_count": int(message_count), "skip_source": bool(skip_source), "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "pst_source_finalizer", PST_SOURCE_KIND, f"{source_rel_path}:__finalize__", source_rel_path, int(commit_order), compact_json_text(ingest_v2_json_safe_value(payload)), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_begin_pst_source_plan( connection: sqlite3.Connection, root: Path, *, run_id: str, rel_path: str, commit_order: int, source_payload: dict[str, object] | None = None, ) -> tuple[dict[str, object] | None, int, bool]: path = ingest_v2_cursor_path(root, rel_path) source_file_size = file_size_bytes(path) source_file_mtime = file_mtime_timestamp(path) pst_payload = dict(source_payload or {}) source_plan_kind = str(pst_payload.get("source_plan_kind") or "pst") source_file_hash = ingest_v2_pst_source_scan_hash( path, {**pst_payload, "source_rel_path": rel_path}, ) existing_source = get_container_source_row(connection, PST_SOURCE_KIND, rel_path) scan_started_at = next_monotonic_utc_timestamp( [ existing_source["last_scan_started_at"] if existing_source is not None else None, existing_source["last_scan_completed_at"] if existing_source is not None else None, ] ) if ( existing_source is not None and container_source_scan_completed(existing_source) and not container_documents_missing_text_revisions( connection, source_kind=PST_SOURCE_KIND, source_rel_path=rel_path, ) and not container_email_documents_missing_threading( connection, source_kind=PST_SOURCE_KIND, source_rel_path=rel_path, ) and existing_source["file_size"] == source_file_size and existing_source["file_hash"] == source_file_hash and (existing_source["file_mtime"] == source_file_mtime or existing_source["file_hash"]) ): message_count = ingest_v2_existing_pst_message_count(connection, source_rel_path=rel_path) ingest_v2_plan_pst_source_finalizer_item( connection, run_id=run_id, source_rel_path=rel_path, source_plan_kind=source_plan_kind, source_file_size=source_file_size, source_file_mtime=source_file_mtime, source_file_hash=source_file_hash, scan_started_at=scan_started_at, message_count=message_count, skip_source=True, commit_order=commit_order, source_action="skipped", ) return None, commit_order + 1, True return ( { "source_rel_path": rel_path, "source_plan_kind": source_plan_kind, "source_action": "new" if existing_source is None else "updated", "source_file_size": source_file_size, "source_file_mtime": source_file_mtime, "source_file_hash": source_file_hash, "scan_started_at": scan_started_at, "next_message_index": 0, "planned_message_count": 0, "next_commit_order": int(commit_order), "message_metadata_by_source_item": dict(pst_payload.get("message_metadata_by_source_item") or {}), "message_match_records": list(pst_payload.get("message_match_records") or []), }, commit_order, False, ) def ingest_v2_plan_current_pst_source( connection: sqlite3.Connection, root: Path, *, run_id: str, current_pst_source: dict[str, object], deadline: float, ) -> tuple[dict[str, object] | None, int, int, bool, dict[str, object]]: source_rel_path = str(current_pst_source["source_rel_path"]) source_plan_kind = str(current_pst_source.get("source_plan_kind") or "pst") path = ingest_v2_cursor_path(root, source_rel_path) next_message_index = int(current_pst_source.get("next_message_index") or 0) next_commit_order = int(current_pst_source.get("next_commit_order") or 1) processed_this_step = 0 plan_started = time.perf_counter() planned_payload_bytes = 0 planned_spilled_payload_bytes = 0 planned_spilled_messages = 0 attachment_count_total = 0 attachment_bytes_total = 0 message_index_start: int | None = None message_index_end: int | None = None reached_end = True for message_index, raw_message in enumerate(iter_pst_messages(path)): if message_index < next_message_index: continue if ( processed_this_step >= INGEST_V2_MBOX_PLAN_BATCH_SIZE or ingest_v2_deadline_remaining_seconds(deadline) < 1.0 ): reached_end = False break raw_message_dict = dict(raw_message) if message_index_start is None: message_index_start = message_index message_index_end = message_index attachment_count, attachment_bytes = ingest_v2_attachment_stats(raw_message_dict.get("attachments")) source_item_id = normalize_source_item_id(raw_message_dict.get("source_item_id")) or f"pst-index:{message_index}" exact_metadata = dict( dict(current_pst_source.get("message_metadata_by_source_item") or {}).get(source_item_id) or {} ) inserted, item_metrics = ingest_v2_plan_pst_message_item( connection, run_id=run_id, source_rel_path=source_rel_path, source_plan_kind=source_plan_kind, message_index=message_index, raw_message=raw_message_dict, source_item_id=source_item_id, source_file_size=( int(current_pst_source["source_file_size"]) if current_pst_source.get("source_file_size") is not None else None ), source_file_mtime=( str(current_pst_source["source_file_mtime"]) if current_pst_source.get("source_file_mtime") is not None else None ), source_file_hash=str(current_pst_source["source_file_hash"]), scan_started_at=str(current_pst_source["scan_started_at"]), commit_order=next_commit_order, message_metadata=exact_metadata, message_match_records=list(current_pst_source.get("message_match_records") or []), ) if inserted: processed_this_step += 1 planned_payload_bytes += int(item_metrics.get("payload_bytes") or 0) planned_spilled_payload_bytes += int(item_metrics.get("spilled_payload_bytes") or 0) planned_spilled_messages += 1 if item_metrics.get("spilled") else 0 attachment_count_total += attachment_count attachment_bytes_total += attachment_bytes next_commit_order += 1 next_message_index = message_index + 1 current_pst_source["next_message_index"] = next_message_index current_pst_source["next_commit_order"] = next_commit_order current_pst_source["planned_message_count"] = ( int(current_pst_source.get("planned_message_count") or 0) + processed_this_step ) plan_metrics = { "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "planned_messages": processed_this_step, "payload_bytes": planned_payload_bytes, "spilled_payload_bytes": planned_spilled_payload_bytes, "spilled_messages": planned_spilled_messages, "attachment_count": attachment_count_total, "attachment_bytes": attachment_bytes_total, "message_index_start": message_index_start, "message_index_end": message_index_end, "duration_ms": ingest_v2_elapsed_ms(plan_started), "source_complete": reached_end, } if not reached_end: return current_pst_source, next_commit_order, processed_this_step, False, plan_metrics ingest_v2_plan_pst_source_finalizer_item( connection, run_id=run_id, source_rel_path=source_rel_path, source_plan_kind=source_plan_kind, source_action=str(current_pst_source.get("source_action") or "updated"), source_file_size=( int(current_pst_source["source_file_size"]) if current_pst_source.get("source_file_size") is not None else None ), source_file_mtime=( str(current_pst_source["source_file_mtime"]) if current_pst_source.get("source_file_mtime") is not None else None ), source_file_hash=str(current_pst_source["source_file_hash"]), scan_started_at=str(current_pst_source["scan_started_at"]), message_count=int(current_pst_source.get("planned_message_count") or 0), skip_source=False, commit_order=next_commit_order, ) return None, next_commit_order + 1, processed_this_step, True, plan_metrics def ingest_v2_production_signature_payload(root: Path, signature: dict[str, object]) -> dict[str, object]: metadata_load_path = Path(signature["metadata_load_path"]) image_load_path = Path(signature["image_load_path"]) if signature.get("image_load_path") is not None else None return { "rel_root": str(signature["rel_root"]), "production_name": str(signature["production_name"]), "metadata_load_rel_path": relative_document_path(root, metadata_load_path), "image_load_rel_path": relative_document_path(root, image_load_path) if image_load_path is not None else None, "source_type": str(signature["source_type"]), } def ingest_v2_production_signature_from_payload(root: Path, payload: dict[str, object]) -> dict[str, object]: image_load_rel_path = normalize_whitespace(str(payload.get("image_load_rel_path") or "")) or None return { "root": (root / str(payload["rel_root"])).resolve(), "rel_root": str(payload["rel_root"]), "production_name": str(payload["production_name"]), "metadata_load_path": root / str(payload["metadata_load_rel_path"]), "image_load_path": (root / image_load_rel_path) if image_load_rel_path else None, "source_type": str(payload["source_type"]), } def ingest_v2_plan_production_row_item( connection: sqlite3.Connection, *, run_id: str, plan: dict[str, object], signature_payload: dict[str, object], commit_order: int, ) -> bool: now = utc_now() production_rel_root = str(plan["production_rel_root"]) control_number = str(plan["control_number"]) rel_path = production_logical_rel_path(production_rel_root, control_number).as_posix() payload = { **plan, "rel_path": rel_path, "metadata_load_rel_path": signature_payload["metadata_load_rel_path"], "image_load_rel_path": signature_payload.get("image_load_rel_path"), "source_type": signature_payload["source_type"], "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "production_row", PRODUCTION_SOURCE_KIND, f"{production_rel_root}:{control_number}", rel_path, int(commit_order), compact_json_text(ingest_v2_json_safe_value(payload)), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_plan_production_preview_batch_item( connection: sqlite3.Connection, *, run_id: str, plan: dict[str, object], signature_payload: dict[str, object], rel_path: str, page_refs: list[dict[str, object]], batch_index: int, parent_order: int, commit_order: int, ) -> bool: if not page_refs: return False now = utc_now() production_rel_root = str(plan["production_rel_root"]) control_number = str(plan["control_number"]) payload = { "production_rel_root": production_rel_root, "production_name": plan["production_name"], "control_number": control_number, "rel_path": rel_path, "metadata_load_rel_path": signature_payload["metadata_load_rel_path"], "image_load_rel_path": signature_payload.get("image_load_rel_path"), "source_type": signature_payload["source_type"], "batch_index": int(batch_index), "page_refs": page_refs, "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, parent_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "production_preview_batch", PRODUCTION_SOURCE_KIND, f"{production_rel_root}:{control_number}:preview:{batch_index}", rel_path, int(commit_order), int(parent_order), compact_json_text(ingest_v2_json_safe_value(payload)), "pending", now, now, ), ) return int(cursor.rowcount or 0) > 0 def ingest_v2_plan_conversation_preview_items( connection: sqlite3.Connection, *, run_id: str, conversation_ids: list[int], next_commit_order: int, ) -> dict[str, int]: now = utc_now() planned = 0 current_order = int(next_commit_order) for conversation_id in sorted(dict.fromkeys(int(value) for value in conversation_ids)): payload = { "conversation_id": conversation_id, "planned_at": now, } cursor = connection.execute( """ INSERT OR IGNORE INTO ingest_work_items ( run_id, unit_type, source_kind, source_key, rel_path, commit_order, payload_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "conversation_preview", "conversation", f"conversation:{conversation_id}", None, current_order, compact_json_text(ingest_v2_json_safe_value(payload)), "pending", now, now, ), ) current_order += 1 if int(cursor.rowcount or 0) > 0: planned += 1 return {"planned": planned, "next_commit_order": current_order} def ingest_v2_conversation_id_from_preview_rel_path(rel_preview_path: object) -> int | None: normalized = normalize_internal_rel_path(Path(str(rel_preview_path or ""))) match = re.match(r"^previews/conversations/conversation-(\d+)(?:/|$)", normalized) if match is None: return None try: return int(match.group(1)) except ValueError: return None def ingest_v2_run_affected_conversation_ids(connection: sqlite3.Connection, *, run_id: str) -> set[int]: rows = connection.execute( """ SELECT affected_document_ids_json FROM ingest_work_items WHERE run_id = ? AND status = 'committed' AND unit_type != 'conversation_preview' AND affected_document_ids_json IS NOT NULL """, (run_id,), ).fetchall() affected_document_ids: set[int] = set() for row in rows: raw_document_ids = decode_json_text(row["affected_document_ids_json"], default=[]) or [] if not isinstance(raw_document_ids, list): continue for value in raw_document_ids: try: affected_document_ids.add(int(value)) except (TypeError, ValueError): continue if not affected_document_ids: return set() placeholders = ", ".join("?" for _ in sorted(affected_document_ids)) conversation_ids = { int(row["conversation_id"]) for row in connection.execute( f""" SELECT DISTINCT conversation_id FROM documents WHERE id IN ({placeholders}) AND conversation_id IS NOT NULL AND lifecycle_status NOT IN ('missing', 'deleted') """, tuple(sorted(affected_document_ids)), ).fetchall() if row["conversation_id"] is not None } preview_rows = connection.execute( f""" SELECT rel_preview_path FROM document_previews WHERE document_id IN ({placeholders}) """, tuple(sorted(affected_document_ids)), ).fetchall() for row in preview_rows: conversation_id = ingest_v2_conversation_id_from_preview_rel_path(row["rel_preview_path"]) if conversation_id is not None: conversation_ids.add(conversation_id) return conversation_ids def ingest_v2_conversation_preview_row_is_stale(document_updated_at: object, preview_created_at: object) -> bool: document_updated = parse_utc_timestamp(document_updated_at) preview_created = parse_utc_timestamp(preview_created_at) if document_updated is None or preview_created is None: return True return preview_created < document_updated def ingest_v2_conversation_previews_need_refresh( connection: sqlite3.Connection, paths: dict[str, Path], conversation_ids: list[int], *, force_conversation_ids: set[int] | None = None, ) -> list[int]: normalized_conversation_ids = sorted(dict.fromkeys(int(value) for value in conversation_ids)) if not normalized_conversation_ids: return [] forced = set(force_conversation_ids or set()) placeholders = ", ".join("?" for _ in normalized_conversation_ids) conversation_rows = connection.execute( f""" SELECT id, conversation_type FROM conversations WHERE id IN ({placeholders}) """, tuple(normalized_conversation_ids), ).fetchall() conversation_type_by_id = { int(row["id"]): normalize_whitespace(str(row["conversation_type"] or "")).lower() for row in conversation_rows } rows = connection.execute( f""" SELECT d.conversation_id, d.id AS document_id, COALESCE(active_tr.created_at, d.updated_at) AS preview_freshness_at, d.content_type, dp.rel_preview_path, dp.preview_type, dp.created_at AS preview_created_at FROM documents d LEFT JOIN text_revisions active_tr ON active_tr.id = COALESCE(d.active_search_text_revision_id, d.source_text_revision_id) LEFT JOIN document_previews dp ON dp.document_id = d.id WHERE d.conversation_id IN ({placeholders}) AND d.lifecycle_status NOT IN ('missing', 'deleted') AND COALESCE(d.child_document_kind, '') != ? ORDER BY d.conversation_id ASC, d.id ASC, dp.ordinal ASC, dp.id ASC """, (*normalized_conversation_ids, CHILD_DOCUMENT_KIND_ATTACHMENT), ).fetchall() doc_ids_by_conversation: dict[int, set[int]] = defaultdict(set) content_types_by_conversation: dict[int, list[str]] = defaultdict(list) preview_rows_by_document: dict[int, list[sqlite3.Row]] = defaultdict(list) preview_freshness_at_by_id: dict[int, object] = {} for row in rows: conversation_id = int(row["conversation_id"]) document_id = int(row["document_id"]) if document_id not in doc_ids_by_conversation[conversation_id]: content_types_by_conversation[conversation_id].append( normalize_whitespace(str(row["content_type"] or "")).lower() ) doc_ids_by_conversation[conversation_id].add(document_id) preview_freshness_at_by_id[document_id] = row["preview_freshness_at"] preview_rows_by_document[document_id].append(row) refresh_ids: set[int] = set() for conversation_id in normalized_conversation_ids: if conversation_id in forced: refresh_ids.add(conversation_id) continue document_ids = doc_ids_by_conversation.get(conversation_id) or set() if not document_ids: continue conversation_type = conversation_type_by_id.get(conversation_id, "") content_types = content_types_by_conversation.get(conversation_id) or [] writes_conversation_artifacts = ( conversation_type == "chat" or (bool(content_types) and all(content_type == "chat" for content_type in content_types)) or len(document_ids) > 1 ) conversation_base = conversation_preview_base_path(conversation_id).as_posix() if writes_conversation_artifacts and not (paths["state_dir"] / conversation_preview_full_rel_path(conversation_id)).exists(): refresh_ids.add(conversation_id) continue stale_or_missing = False for document_id in sorted(document_ids): non_native_rows = [ row for row in preview_rows_by_document.get(document_id, []) if row["rel_preview_path"] is not None and normalize_whitespace(str(row["preview_type"] or "")).lower() != "native" ] if not non_native_rows: stale_or_missing = True break if writes_conversation_artifacts and not any( normalize_internal_rel_path(Path(str(row["rel_preview_path"] or ""))).startswith(conversation_base + "/") for row in non_native_rows ): stale_or_missing = True break for row in non_native_rows: rel_preview_path = normalize_whitespace(str(row["rel_preview_path"] or "")) if rel_preview_path and not (paths["state_dir"] / rel_preview_path).exists(): stale_or_missing = True break if ingest_v2_conversation_preview_row_is_stale( preview_freshness_at_by_id.get(document_id), row["preview_created_at"], ): stale_or_missing = True break if stale_or_missing: break if stale_or_missing: refresh_ids.add(conversation_id) continue if writes_conversation_artifacts: orphan_rows = connection.execute( """ SELECT dp.document_id, d.conversation_id, d.lifecycle_status FROM document_previews dp LEFT JOIN documents d ON d.id = dp.document_id WHERE dp.rel_preview_path LIKE ? """, (conversation_base + "/%",), ).fetchall() for row in orphan_rows: if ( row["conversation_id"] is None or int(row["conversation_id"]) != conversation_id or normalize_whitespace(str(row["lifecycle_status"] or "")).lower() in {"missing", "deleted"} ): refresh_ids.add(conversation_id) break return [ conversation_id for conversation_id in normalized_conversation_ids if conversation_id in refresh_ids ] def ingest_v2_plan_production_root( connection: sqlite3.Connection, root: Path, *, run_id: str, signature_payload: dict[str, object], next_commit_order: int, paths: dict[str, Path] | None = None, ) -> dict[str, object]: resolved_paths = paths if paths is not None else workspace_paths(root) signature = ingest_v2_production_signature_from_payload(root, signature_payload) production_root = Path(signature["root"]) metadata_load_path = Path(signature["metadata_load_path"]) image_load_path = Path(signature["image_load_path"]) if signature.get("image_load_path") is not None else None metadata = parse_production_metadata_load(metadata_load_path) image_rows = parse_production_image_load(image_load_path) resolved_image_rows: list[dict[str, object]] = [] for image_row in image_rows: resolved_path = resolve_production_source_path(root, production_root, image_row["image_path"]) resolved_image_rows.append({**image_row, "resolved_path": resolved_path}) production_row_plans, seen_control_numbers = plan_production_record_work( root, production_root, signature, list(metadata["rows"]), resolved_image_rows, ) production_row = connection.execute( """ SELECT id FROM productions WHERE rel_root = ? """, (str(signature["rel_root"]),), ).fetchone() existing_production_id = int(production_row["id"]) if production_row is not None else None existing_by_control_number: dict[str, sqlite3.Row] = {} if existing_production_id is not None: existing_rows = connection.execute( """ SELECT * FROM documents WHERE production_id = ? AND control_number IS NOT NULL """, (existing_production_id,), ).fetchall() existing_by_control_number = { str(row["control_number"]): row for row in existing_rows if row["control_number"] } planned_rows = 0 planned_preview_batches = 0 skipped_unchanged_rows = 0 current_order = int(next_commit_order) for plan in production_row_plans: production_rel_root = str(plan["production_rel_root"]) control_number = str(plan["control_number"]) rel_path = production_logical_rel_path(production_rel_root, control_number).as_posix() if ( existing_production_id is not None and production_plan_matches_existing_document( resolved_paths, connection, root, existing_by_control_number.get(control_number), plan, production_id=existing_production_id, preview_image_max_dimension=( INGEST_V2_PRODUCTION_PREVIEW_IMAGE_MAX_DIMENSION if INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES else None ), embed_preview_images=INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES, ) ): skipped_unchanged_rows += 1 continue row_order = current_order inserted = ingest_v2_plan_production_row_item( connection, run_id=run_id, plan=plan, signature_payload=signature_payload, commit_order=row_order, ) current_order += 1 if inserted: planned_rows += 1 native_path = Path(str(plan["native_path"])) if plan.get("native_path") is not None else None page_refs = ( production_preview_page_asset_refs( rel_path, control_number, [ Path(str(path)) for path in list(plan.get("matching_image_paths") or []) if path ], ) if ( not INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES and production_previewable_native(native_path) is None ) else [] ) for batch_index, start in enumerate(range(0, len(page_refs), INGEST_V2_PRODUCTION_PREVIEW_BATCH_SIZE), start=1): batch_refs = page_refs[start : start + INGEST_V2_PRODUCTION_PREVIEW_BATCH_SIZE] batch_inserted = ingest_v2_plan_production_preview_batch_item( connection, run_id=run_id, plan=plan, signature_payload=signature_payload, rel_path=rel_path, page_refs=batch_refs, batch_index=batch_index, parent_order=row_order, commit_order=current_order, ) current_order += 1 if batch_inserted: planned_preview_batches += 1 return { **signature_payload, "planned_rows": planned_rows, "planned_preview_batches": planned_preview_batches, "skipped_unchanged_rows": skipped_unchanged_rows, "seen_control_numbers": sorted(seen_control_numbers), "docs_missing_linked_text": sum(int(plan["missing_linked_text"]) for plan in production_row_plans), "docs_missing_linked_images": sum(int(plan["missing_linked_images"]) for plan in production_row_plans), "docs_missing_linked_natives": sum(int(plan["missing_linked_natives"]) for plan in production_row_plans), "next_commit_order": current_order, } def ingest_v2_planning_child_paths(root: Path, directory: Path, *, recursive: bool) -> list[str]: child_paths: list[str] = [] for child in sorted(directory.iterdir(), key=lambda item: item.name): if not path_is_at_or_under(child, root): continue rel_path = ingest_v2_cursor_rel_path(root, child) if ingest_v2_path_is_workspace_internal(root, child): continue if child.is_dir() and not recursive: continue child_paths.append(rel_path) return child_paths def ingest_v2_worker_id(prefix: str) -> str: return f"{prefix}-{os.getpid()}-{secrets.token_hex(4)}" def ingest_v2_json_safe_value(value: object) -> object: if isinstance(value, bytes): return {INGEST_V2_BYTES_B64_KEY: base64.b64encode(value).decode("ascii")} if isinstance(value, Path): return str(value) if isinstance(value, sqlite3.Row): return {key: ingest_v2_json_safe_value(value[key]) for key in value.keys()} if isinstance(value, dict): return {str(key): ingest_v2_json_safe_value(child_value) for key, child_value in value.items()} if isinstance(value, (list, tuple)): return [ingest_v2_json_safe_value(child_value) for child_value in value] if value is None or isinstance(value, (str, int, float, bool)): return value return str(value) def ingest_v2_json_restore_value(value: object) -> object: if isinstance(value, dict): if set(value.keys()) == {INGEST_V2_BYTES_B64_KEY}: try: return base64.b64decode(str(value[INGEST_V2_BYTES_B64_KEY]).encode("ascii")) except Exception: return b"" return {str(key): ingest_v2_json_restore_value(child_value) for key, child_value in value.items()} if isinstance(value, list): return [ingest_v2_json_restore_value(child_value) for child_value in value] return value def ingest_v2_clone_cursor(cursor: dict[str, object]) -> dict[str, object]: cloned = decode_json_text(compact_json_text(ingest_v2_json_safe_value(cursor)), default={}) or {} return dict(cloned) if isinstance(cloned, dict) else {} def ingest_v2_deadline_remaining_seconds(deadline: float) -> float: return max(0.0, deadline - time.perf_counter()) def ingest_v2_stale_lease_cutoff(now_dt: datetime | None = None) -> str: return format_utc_timestamp( (now_dt or datetime.now(timezone.utc)) - timedelta(seconds=INGEST_V2_WORK_ITEM_LEASE_SECONDS) ) def ingest_v2_lease_is_active( expires_at: object, refreshed_at: object, *, now: datetime | None = None, ) -> bool: now_dt = now or datetime.now(timezone.utc) parsed_refresh = parse_utc_timestamp(refreshed_at) if parsed_refresh is None: return False if parsed_refresh <= now_dt - timedelta(seconds=INGEST_V2_WORK_ITEM_LEASE_SECONDS): return False return lease_is_active(expires_at, now=now_dt) def ingest_v2_reclaim_stale_prepare_items(connection: sqlite3.Connection, *, run_id: str) -> int: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) stale_cutoff = ingest_v2_stale_lease_cutoff(now_dt) connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, updated_at = ? WHERE run_id = ? AND status = 'leased' AND lease_expires_at IS NOT NULL AND ( lease_expires_at <= ? OR updated_at <= ? ) """, (now, run_id, now, stale_cutoff), ) reclaimed = int(cursor.rowcount or 0) if reclaimed: connection.execute( """ INSERT INTO ingest_worker_events ( run_id, worker_id, event_type, phase, details_json, created_at ) VALUES (?, ?, ?, ?, ?, ?) """, ( run_id, None, "reclaim_stale_prepare_items", "prepare", compact_json_text({"reclaimed": reclaimed}), now, ), ) connection.commit() return reclaimed except Exception: connection.rollback() raise def ingest_v2_claim_prepare_items( connection: sqlite3.Connection, *, run_id: str, worker_id: str, limit: int = INGEST_V2_PREPARE_BATCH_SIZE, ) -> tuple[list[sqlite3.Row], bool]: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) stale_cutoff = ingest_v2_stale_lease_cutoff(now_dt) lease_expires_at = lease_expiration_after(INGEST_V2_WORK_ITEM_LEASE_SECONDS, now=now_dt) connection.execute("BEGIN IMMEDIATE") try: row = require_ingest_v2_run_row(connection, run_id) if ( str(row["phase"]) != "preparing" or str(row["status"]) in INGEST_V2_TERMINAL_STATUSES or row["cancel_requested_at"] is not None ): connection.rollback() return [], False active_workers = int( connection.execute( """ SELECT COUNT(DISTINCT lease_owner) FROM ingest_work_items WHERE run_id = ? AND status = 'leased' AND lease_owner IS NOT NULL AND lease_owner != ? AND lease_expires_at IS NOT NULL AND lease_expires_at > ? AND updated_at > ? """, (run_id, worker_id, now, stale_cutoff), ).fetchone()[0] or 0 ) soft_limit = int(row["prepare_worker_soft_limit"] or DEFAULT_WORKER_BACKGROUND_MAX_PARALLEL) if active_workers >= soft_limit: connection.execute( """ UPDATE ingest_runs SET last_heartbeat_at = ? WHERE run_id = ? """, (now, run_id), ) connection.commit() return [], True claim_rows = connection.execute( """ SELECT * FROM ingest_work_items WHERE run_id = ? AND unit_type IN ( 'loose_file', 'production_row', 'production_preview_batch', 'slack_conversation', 'slack_document', 'conversation_preview', 'mbox_message', 'mbox_source_finalizer', 'pst_message', 'pst_source_finalizer' ) AND status = 'pending' ORDER BY commit_order ASC, id ASC LIMIT ? """, (run_id, max(1, int(limit))), ).fetchall() claim_ids = [int(claim_row["id"]) for claim_row in claim_rows] if claim_ids: placeholders = ",".join("?" for _ in claim_ids) connection.execute( f""" UPDATE ingest_work_items SET status = 'leased', lease_owner = ?, lease_expires_at = ?, attempts = attempts + 1, updated_at = ? WHERE run_id = ? AND status = 'pending' AND id IN ({placeholders}) """, (worker_id, lease_expires_at, now, run_id, *claim_ids), ) connection.execute( """ UPDATE ingest_runs SET last_heartbeat_at = ? WHERE run_id = ? """, (now, run_id), ) connection.commit() except Exception: connection.rollback() raise if not claim_ids: return [], False placeholders = ",".join("?" for _ in claim_ids) claimed_rows = connection.execute( f""" SELECT * FROM ingest_work_items WHERE run_id = ? AND lease_owner = ? AND status = 'leased' AND id IN ({placeholders}) ORDER BY commit_order ASC, id ASC """, (run_id, worker_id, *claim_ids), ).fetchall() return list(claimed_rows), False def ingest_v2_release_prepare_item( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, worker_id: str, reason: str, ) -> bool: now = utc_now() connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, updated_at = ?, last_error = ? WHERE run_id = ? AND id = ? AND status = 'leased' AND lease_owner = ? """, (now, reason, run_id, work_item_id, worker_id), ) released = int(cursor.rowcount or 0) > 0 connection.commit() return released except Exception: connection.rollback() raise def ingest_v2_prepare_claim_limit() -> int: return max( INGEST_V2_PREPARE_BATCH_SIZE, ingest_prepare_worker_count(), ingest_container_prepare_worker_count(), ) def ingest_v2_mark_prepare_deferred_timeout( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, worker_id: str, message: str, ) -> bool: now = utc_now() connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'deferred_timeout', lease_owner = NULL, lease_expires_at = NULL, updated_at = ?, last_error = ? WHERE run_id = ? AND id = ? AND status = 'leased' AND lease_owner = ? """, (now, message, run_id, work_item_id, worker_id), ) marked = int(cursor.rowcount or 0) > 0 if marked: connection.execute( """ INSERT INTO ingest_worker_events ( run_id, worker_id, event_type, work_item_id, phase, details_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( run_id, worker_id, "prepare_deferred_timeout", work_item_id, "prepare", compact_json_text({"message": message}), now, ), ) connection.commit() return marked except Exception: connection.rollback() raise def ingest_v2_prepare_worker_count_for_rows(rows: list[sqlite3.Row]) -> int: if not rows: return 1 unit_types = {str(row["unit_type"] or "") for row in rows} container_unit_types = { "mbox_message", "mbox_source_finalizer", "pst_message", "pst_source_finalizer", } if unit_types and unit_types <= container_unit_types: return ingest_container_prepare_worker_count() if unit_types.isdisjoint(container_unit_types): return ingest_prepare_worker_count() return max(ingest_prepare_worker_count(), ingest_container_prepare_worker_count()) def ingest_v2_prepare_claimed_work_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, ) -> dict[str, object]: work_item_id = int(work_item_row["id"]) if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return { "work_item_id": work_item_id, "payload_kind": str(work_item_row["unit_type"] or "loose_file"), "prepared_item": None, "source_fingerprint": {}, "defer_message": "Not enough budget remaining to start prepare.", } unit_type = str(work_item_row["unit_type"] or "") if unit_type == "production_row": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_production_row_item( root, work_item_row, deadline=deadline, ) payload_kind = "production_row" elif unit_type == "production_preview_batch": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_production_preview_batch_item( root, work_item_row, deadline=deadline, ) payload_kind = "production_preview_batch" elif unit_type == "slack_conversation": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_slack_conversation_item( work_item_row, deadline=deadline, ) payload_kind = "slack_conversation" elif unit_type == "slack_document": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_slack_document_item( work_item_row, deadline=deadline, ) payload_kind = "slack_document" elif unit_type == "conversation_preview": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_conversation_preview_item( work_item_row, deadline=deadline, ) payload_kind = "conversation_preview" elif unit_type == "mbox_message": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_mbox_message_item( root, work_item_row, deadline=deadline, ) payload_kind = "mbox_message" elif unit_type == "mbox_source_finalizer": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_mbox_source_finalizer_item( work_item_row, ) payload_kind = "mbox_source_finalizer" elif unit_type == "pst_message": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_pst_message_item( root, work_item_row, deadline=deadline, ) payload_kind = "pst_message" elif unit_type == "pst_source_finalizer": prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_pst_source_finalizer_item( work_item_row, ) payload_kind = "pst_source_finalizer" else: prepared_item, source_fingerprint, defer_message = ingest_v2_prepare_loose_file_item( root, work_item_row, deadline=deadline, ) payload_kind = "loose_file" return { "work_item_id": work_item_id, "payload_kind": payload_kind, "prepared_item": prepared_item, "source_fingerprint": source_fingerprint, "defer_message": defer_message, } def ingest_v2_prepare_claimed_work_items_parallel( root: Path, rows: list[sqlite3.Row], *, deadline: float, prepare_workers: int, ) -> list[dict[str, object]]: if not rows: return [] effective_workers = max(1, min(int(prepare_workers), len(rows))) if effective_workers == 1: return [ ingest_v2_prepare_claimed_work_item(root, row, deadline=deadline) for row in rows ] with ThreadPoolExecutor(max_workers=effective_workers) as executor: futures = [ executor.submit( ingest_v2_prepare_claimed_work_item, root, row, deadline=deadline, ) for row in rows ] return [future.result() for future in futures] def ingest_v2_prepare_loose_file_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} payload_dict = payload if isinstance(payload, dict) else {} rel_path = str(work_item_row["rel_path"] or payload_dict.get("rel_path") or "") path = ingest_v2_cursor_path(root, rel_path) file_size, file_mtime_ns = source_file_snapshot(path) source_fingerprint = { "rel_path": rel_path, "size": file_size, "mtime_ns": file_mtime_ns, "hash": None, } if file_size is not None and file_size > INGEST_V2_MAX_SINGLE_STEP_HASH_BYTES: return ( None, source_fingerprint, ( f"File is too large for one bounded prepare step ({file_size} bytes > " f"{INGEST_V2_MAX_SINGLE_STEP_HASH_BYTES} bytes)." ), ) if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." file_type = str(payload_dict.get("file_type") or normalize_extension(path)) item = { "rel_path": rel_path, "path": str(path), "file_type": file_type, "source_file_size": file_size, "source_file_mtime_ns": file_mtime_ns, } hash_started = time.perf_counter() item["file_hash"] = sha256_file(path) hash_ms = ingest_v2_elapsed_ms(hash_started) source_fingerprint["hash"] = item["file_hash"] prepared_item = prepare_loose_file_item(item) prepared_item["prepare_hash_ms"] = hash_ms return prepared_item, source_fingerprint, None def ingest_v2_prepare_production_row_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} payload_dict = payload if isinstance(payload, dict) else {} production_rel_root = str(payload_dict.get("production_rel_root") or "") control_number = str(payload_dict.get("control_number") or "") source_fingerprint = { "production_rel_root": production_rel_root, "control_number": control_number, "metadata_load_rel_path": payload_dict.get("metadata_load_rel_path"), "image_load_rel_path": payload_dict.get("image_load_rel_path"), } if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." rel_path = str(payload_dict.get("rel_path") or production_logical_rel_path(production_rel_root, control_number)) matching_image_paths = [ Path(str(path)) for path in list(payload_dict.get("matching_image_paths") or []) if path ] preview_image_refs = production_preview_page_asset_refs(rel_path, control_number, matching_image_paths) prepared_item = prepare_production_row_plan( root, payload_dict, preview_image_max_dimension=( INGEST_V2_PRODUCTION_PREVIEW_IMAGE_MAX_DIMENSION if INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES else None ), preview_image_refs=preview_image_refs, embed_preview_images=INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES, ) prepared_item["prepare_hash_ms"] = 0.0 return prepared_item, source_fingerprint, None def ingest_v2_prepare_production_preview_batch_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} payload_dict = payload if isinstance(payload, dict) else {} page_refs = [ dict(ref) for ref in list(payload_dict.get("page_refs") or []) if isinstance(ref, dict) ] source_fingerprint = { "production_rel_root": payload_dict.get("production_rel_root"), "control_number": payload_dict.get("control_number"), "batch_index": payload_dict.get("batch_index"), "page_count": len(page_refs), } if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." prepare_started = time.perf_counter() if INGEST_V2_PRODUCTION_PREVIEW_EMBED_IMAGES: prepared_item = { **payload_dict, "payload_kind": "production_preview_batch", "source_kind": PRODUCTION_SOURCE_KIND, "page_assets": [], "prepare_ms": ingest_v2_elapsed_ms(prepare_started), "prepare_hash_ms": 0.0, "prepare_extract_ms": ingest_v2_elapsed_ms(prepare_started), "prepare_chunk_ms": 0.0, "prepare_error": None, } return prepared_item, source_fingerprint, None page_assets: list[dict[str, object]] = [] for ref in page_refs: if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." source_path = Path(str(ref.get("source_path") or "")) if not source_path.exists(): continue preview_raster = image_path_preview_raster( source_path, max_dimension=INGEST_V2_PRODUCTION_PREVIEW_IMAGE_MAX_DIMENSION, ) if preview_raster is None: continue preview_bytes, _, _ = preview_raster page_assets.append( { "ordinal": int(ref.get("ordinal") or 0), "label": str(ref.get("label") or ""), "rel_preview_path": str(ref.get("rel_preview_path") or ""), "payload": preview_bytes, } ) prepared_item = { **payload_dict, "payload_kind": "production_preview_batch", "source_kind": PRODUCTION_SOURCE_KIND, "page_assets": page_assets, "prepare_ms": ingest_v2_elapsed_ms(prepare_started), "prepare_hash_ms": 0.0, "prepare_extract_ms": ingest_v2_elapsed_ms(prepare_started), "prepare_chunk_ms": 0.0, "prepare_error": None, } return prepared_item, source_fingerprint, None def ingest_v2_prepare_conversation_preview_item( work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} payload_dict = payload if isinstance(payload, dict) else {} conversation_id = int(payload_dict.get("conversation_id") or 0) source_fingerprint = {"conversation_id": conversation_id} if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." return ( { "payload_kind": "conversation_preview", "conversation_id": conversation_id, "prepare_ms": 0.0, "prepare_hash_ms": 0.0, "prepare_extract_ms": 0.0, "prepare_chunk_ms": 0.0, "prepare_error": None, }, source_fingerprint, None, ) def ingest_v2_prepare_slack_conversation_item( work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} payload_dict = payload if isinstance(payload, dict) else {} rel_paths = [str(rel_path) for rel_path in list(payload_dict.get("rel_paths") or [])] conversation_identity = payload_dict.get("conversation_identity") identity_root = "" if isinstance(conversation_identity, list) and len(conversation_identity) > 1: identity_root = str(conversation_identity[1]) source_locator = str(payload_dict.get("source_locator") or identity_root) source_fingerprint = { "source_locator": source_locator, "conversation_key": payload_dict.get("conversation_key"), "rel_paths": rel_paths, } if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." prepared_item = prepare_slack_conversation_plan(payload_dict) prepared_item["payload_kind"] = "slack_conversation" prepared_item["source_kind"] = SLACK_EXPORT_SOURCE_KIND prepared_item["source_locator"] = source_locator prepared_item["prepare_hash_ms"] = 0.0 prepared_item["prepare_extract_ms"] = max( 0.0, float(prepared_item.get("prepare_ms") or 0.0) - float(prepared_item.get("prepare_chunk_ms") or 0.0), ) return prepared_item, source_fingerprint, None def ingest_v2_prepare_slack_document_item( work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} payload_dict = payload if isinstance(payload, dict) else {} plan = dict(payload_dict.get("plan") or {}) rel_path = str(plan.get("rel_path") or "") source_locator = str(payload_dict.get("source_locator") or "") source_fingerprint = { "source_locator": source_locator, "conversation_key": payload_dict.get("conversation_key"), "rel_path": rel_path, } if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." prepared_item = prepare_slack_document_plan(payload_dict) prepared_item["payload_kind"] = "slack_document" prepared_item["source_kind"] = SLACK_EXPORT_SOURCE_KIND prepared_item["source_locator"] = source_locator prepared_item["prepare_hash_ms"] = 0.0 prepared_item["prepare_extract_ms"] = max( 0.0, float(prepared_item.get("prepare_ms") or 0.0) - float(prepared_item.get("prepare_chunk_ms") or 0.0), ) return prepared_item, source_fingerprint, None def ingest_v2_mbox_message_from_payload(root: Path, payload: dict[str, object]) -> dict[str, object]: source_rel_path = str(payload["source_rel_path"]) source_item_id = str(payload["source_item_id"]) inline_payload = ingest_v2_json_restore_value(payload.get("payload_bytes")) if isinstance(inline_payload, bytes) and inline_payload: return { "source_item_id": source_item_id, "payload_hash": normalize_whitespace(str(payload.get("payload_hash") or "")) or sha256_bytes(inline_payload), "parsed_message": BytesParser(policy=policy.default).parsebytes(inline_payload), } path = ingest_v2_cursor_path(root, source_rel_path) message_key = payload.get("message_key") message_index = int(payload.get("message_index") or 0) archive = mailbox.mbox(str(path), factory=mailbox.mboxMessage, create=False) try: raw_message = None if message_key is not None: try: raw_message = archive.get_message(message_key) except Exception: try: raw_message = archive.get_message(int(message_key)) except Exception: raw_message = None if raw_message is None: for index, candidate in enumerate(archive): if index == message_index: raw_message = candidate break if raw_message is None: raise RetrieverError(f"MBOX message {message_index} is missing from {source_rel_path}.") payload_bytes = raw_message.as_bytes(policy=policy.default, unixfrom=False) payload_hash = sha256_bytes(payload_bytes) parsed_message = BytesParser(policy=policy.default).parsebytes(payload_bytes) return { "source_item_id": source_item_id, "payload_hash": payload_hash, "parsed_message": parsed_message, } finally: try: archive.close() except Exception: pass def ingest_v2_pst_message_from_payload(payload: dict[str, object]) -> dict[str, object]: spill_rel_path = normalize_whitespace(str(payload.get("raw_message_spill_rel_path") or "")) or None if spill_rel_path is not None: active_paths = active_workspace_paths() if active_paths is None: raise RetrieverError("Cannot hydrate PST work item payload without an active workspace.") spilled_payload = ingest_v2_load_spilled_payload(active_paths["state_dir"] / spill_rel_path) raw_message = spilled_payload.get("raw_message") else: raw_message = ingest_v2_json_restore_value(payload.get("raw_message") or {}) if not isinstance(raw_message, dict): raise RetrieverError("PST work item is missing a raw message payload.") restored = dict(raw_message) restored["source_item_id"] = str(payload.get("source_item_id") or restored.get("source_item_id") or "") return restored def ingest_v2_container_source_config(source_kind: str) -> dict[str, object]: if source_kind == MBOX_SOURCE_KIND: return { "source_label": "MBOX", "source_plan_kind_default": "mbox", "message_payload_kind": "mbox_message", "source_finalizer_payload_kind": "mbox_source_finalizer", "message_commit_event_type": "commit_mbox_message", "source_stats_key": "mbox_stats", "created_stat_key": "mbox_messages_created", "updated_stat_key": "mbox_messages_updated", "skipped_stat_key": "mbox_sources_skipped", "finalized_stat_key": "mbox_sources_finalized", "deleted_stat_key": "mbox_messages_deleted", "commit_batch_max_items": INGEST_V2_MBOX_COMMIT_BATCH_MAX_ITEMS, "commit_batch_max_payload_bytes": INGEST_V2_MBOX_COMMIT_BATCH_MAX_PAYLOAD_BYTES, "dataset_name_builder": mbox_dataset_name, } if source_kind == PST_SOURCE_KIND: return { "source_label": "PST", "source_plan_kind_default": "pst", "message_payload_kind": "pst_message", "source_finalizer_payload_kind": "pst_source_finalizer", "message_commit_event_type": "commit_pst_message", "source_stats_key": "pst_stats", "created_stat_key": "pst_messages_created", "updated_stat_key": "pst_messages_updated", "skipped_stat_key": "pst_sources_skipped", "finalized_stat_key": "pst_sources_finalized", "deleted_stat_key": "pst_messages_deleted", "commit_batch_max_items": INGEST_V2_PST_COMMIT_BATCH_MAX_ITEMS, "commit_batch_max_payload_bytes": INGEST_V2_PST_COMMIT_BATCH_MAX_PAYLOAD_BYTES, "dataset_name_builder": pst_dataset_name, } raise RetrieverError(f"Unsupported container source kind: {source_kind!r}") def ingest_v2_work_item_payload_dict(work_item_row: sqlite3.Row) -> dict[str, object]: payload = decode_json_text(work_item_row["payload_json"], default={}) or {} return payload if isinstance(payload, dict) else {} def ingest_v2_build_container_source_fingerprint( root: Path, payload_dict: dict[str, object], *, include_message_key: bool = False, ) -> dict[str, object]: source_rel_path = str(payload_dict.get("source_rel_path") or "") path = ingest_v2_cursor_path(root, source_rel_path) source_fingerprint = { "source_rel_path": source_rel_path, "message_index": payload_dict.get("message_index"), "source_item_id": payload_dict.get("source_item_id"), "size": file_size_bytes(path), "mtime": file_mtime_timestamp(path), "hash": payload_dict.get("source_file_hash"), } if include_message_key: source_fingerprint["message_key"] = payload_dict.get("message_key") return source_fingerprint def ingest_v2_prepare_source_changed_item( source_label: str, source_rel_path: str, source_item_id: object, ) -> dict[str, object]: return { "prepare_error": f"{source_label} source changed after planning: {source_rel_path}", "prepare_ms": 0.0, "prepare_hash_ms": 0.0, "prepare_extract_ms": 0.0, "prepare_chunk_ms": 0.0, "source_rel_path": source_rel_path, "source_item_id": str(source_item_id or ""), } def ingest_v2_finalize_prepared_container_message_item( prepared_item: dict[str, object], payload_dict: dict[str, object], *, source_kind: str, source_plan_kind: str, source_rel_path: str, prepare_started: float, ) -> dict[str, object]: prepared_item["source_kind"] = source_kind prepared_item["source_plan_kind"] = source_plan_kind prepared_item["source_rel_path"] = source_rel_path prepared_item["scan_started_at"] = str(payload_dict["scan_started_at"]) prepared_item["source_file_size"] = payload_dict.get("source_file_size") prepared_item["source_file_mtime"] = payload_dict.get("source_file_mtime") prepared_item["source_file_hash"] = payload_dict.get("source_file_hash") prepared_item["prepare_hash_ms"] = 0.0 prepared_item["prepare_extract_ms"] = max( 0.0, float(prepared_item.get("prepare_ms") or 0.0) - float(prepared_item.get("prepare_chunk_ms") or 0.0), ) prepared_item["prepare_total_step_ms"] = ingest_v2_elapsed_ms(prepare_started) return prepared_item def ingest_v2_build_mbox_message_normalizer( root: Path, payload_dict: dict[str, object], ) -> Callable[[str, dict[str, object]], dict[str, object] | None]: source_plan_kind = str(payload_dict.get("source_plan_kind") or "mbox") if source_plan_kind != "gmail": return normalize_mbox_message def normalize_v2_gmail_message( source_rel_path_for_message: str, message_dict: dict[str, object], ) -> dict[str, object] | None: normalized = normalize_mbox_message(source_rel_path_for_message, message_dict) message_metadata = dict(payload_dict.get("message_metadata") or {}) linked_drive_records = [ dict(record) for record in list(payload_dict.get("linked_drive_records") or []) if isinstance(record, dict) ] linked_drive_attachment_records = [ ingest_v2_gmail_drive_record_from_payload(root, dict(record)) for record in list(payload_dict.get("linked_drive_attachment_records") or []) if isinstance(record, dict) ] extracted = apply_gmail_email_export_metadata( dict(normalized["extracted"]), message_metadata=message_metadata, linked_drive_records=linked_drive_records, ) attachment_payloads = [ attachment for attachment in ( gmail_drive_attachment_payload(dict(record)) for record in linked_drive_attachment_records ) if attachment is not None ] if attachment_payloads: extracted["attachments"] = [*list(extracted.get("attachments") or []), *attachment_payloads] normalized["extracted"] = extracted normalized["file_hash"] = gmail_enriched_message_file_hash( normalized.get("file_hash"), message_metadata=message_metadata, linked_drive_records=linked_drive_records, linked_drive_attachment_records=linked_drive_attachment_records, ) return normalized return normalize_v2_gmail_message def ingest_v2_build_pst_message_normalizer( payload_dict: dict[str, object], ) -> Callable[[str, dict[str, object]], dict[str, object] | None]: exact_metadata_by_source_item = { str(key): dict(value) for key, value in dict(payload_dict.get("message_metadata_by_source_item") or {}).items() if isinstance(value, dict) } message_match_records = [ dict(record) for record in list(payload_dict.get("message_match_records") or []) if isinstance(record, dict) ] def normalize_v2_pst_message( source_rel_path_for_message: str, message_dict: dict[str, object], ) -> dict[str, object] | None: normalized = normalize_pst_message(source_rel_path_for_message, message_dict) if normalized is None: return None message_metadata = select_pst_export_message_metadata( normalized, exact_metadata_by_source_item=exact_metadata_by_source_item, message_match_records=message_match_records, ) if not message_metadata: return normalized enriched = dict(normalized) enriched["extracted"] = apply_pst_export_message_metadata( dict(normalized["extracted"]), message_metadata=message_metadata, identifier_scope=source_rel_path_for_message, ) enriched["file_hash"] = pst_export_enriched_message_file_hash( normalized.get("file_hash"), message_metadata=message_metadata, ) return enriched return normalize_v2_pst_message def ingest_v2_prepare_container_message_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, source_kind: str, raw_message_loader: Callable[[dict[str, object]], dict[str, object]], normalize_message_builder: Callable[ [dict[str, object]], Callable[[str, dict[str, object]], dict[str, object] | None], ], include_message_key: bool = False, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: container_config = ingest_v2_container_source_config(source_kind) payload_dict = ingest_v2_work_item_payload_dict(work_item_row) source_rel_path = str(payload_dict.get("source_rel_path") or "") source_fingerprint = ingest_v2_build_container_source_fingerprint( root, payload_dict, include_message_key=include_message_key, ) if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_PREPARE_MIN_START_SECONDS: return None, source_fingerprint, "Not enough budget remaining to start prepare." if ( payload_dict.get("source_file_size") is not None and int(payload_dict.get("source_file_size") or 0) != int(source_fingerprint["size"] or 0) ): return ( ingest_v2_prepare_source_changed_item( str(container_config["source_label"]), source_rel_path, payload_dict.get("source_item_id"), ), source_fingerprint, None, ) prepare_started = time.perf_counter() prepared_item = prepare_container_message_item( source_rel_path, raw_message_loader(payload_dict), normalize_message_builder(payload_dict), ) source_plan_kind = str( payload_dict.get("source_plan_kind") or container_config["source_plan_kind_default"] ) return ( ingest_v2_finalize_prepared_container_message_item( prepared_item, payload_dict, source_kind=source_kind, source_plan_kind=source_plan_kind, source_rel_path=source_rel_path, prepare_started=prepare_started, ), source_fingerprint, None, ) def ingest_v2_prepare_mbox_message_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: return ingest_v2_prepare_container_message_item( root, work_item_row, deadline=deadline, source_kind=MBOX_SOURCE_KIND, raw_message_loader=lambda payload_dict: ingest_v2_mbox_message_from_payload(root, payload_dict), normalize_message_builder=lambda payload_dict: ingest_v2_build_mbox_message_normalizer(root, payload_dict), include_message_key=True, ) def ingest_v2_prepare_pst_message_item( root: Path, work_item_row: sqlite3.Row, *, deadline: float, ) -> tuple[dict[str, object] | None, dict[str, object], str | None]: return ingest_v2_prepare_container_message_item( root, work_item_row, deadline=deadline, source_kind=PST_SOURCE_KIND, raw_message_loader=ingest_v2_pst_message_from_payload, normalize_message_builder=ingest_v2_build_pst_message_normalizer, ) def ingest_v2_prepare_container_source_finalizer_item_common( work_item_row: sqlite3.Row, *, source_kind: str, payload_kind: str | None = None, ) -> tuple[dict[str, object], dict[str, object], None]: payload_dict = ingest_v2_work_item_payload_dict(work_item_row) prepared_item = { **payload_dict, "source_kind": source_kind, "prepare_ms": 0.0, "prepare_hash_ms": 0.0, "prepare_extract_ms": 0.0, "prepare_chunk_ms": 0.0, "prepare_error": None, } if payload_kind is not None: prepared_item["payload_kind"] = payload_kind source_fingerprint = { "source_rel_path": payload_dict.get("source_rel_path"), "size": payload_dict.get("source_file_size"), "mtime": payload_dict.get("source_file_mtime"), "hash": payload_dict.get("source_file_hash"), "message_count": payload_dict.get("message_count"), } return prepared_item, source_fingerprint, None def ingest_v2_prepare_mbox_source_finalizer_item( work_item_row: sqlite3.Row, ) -> tuple[dict[str, object], dict[str, object], None]: return ingest_v2_prepare_container_source_finalizer_item_common( work_item_row, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_prepare_pst_source_finalizer_item( work_item_row: sqlite3.Row, ) -> tuple[dict[str, object], dict[str, object], None]: return ingest_v2_prepare_container_source_finalizer_item_common( work_item_row, source_kind=PST_SOURCE_KIND, payload_kind="pst_source_finalizer", ) def ingest_v2_hydrate_prepared_production_item(prepared_item: dict[str, object]) -> dict[str, object]: hydrated = dict(prepared_item) for key in ( "text_path", "native_path", "available_text_path", "available_native_path", "preferred_native", "preferred_source_path", ): value = hydrated.get(key) if value: hydrated[key] = Path(str(value)) hydrated["matching_image_paths"] = [ Path(str(path)) for path in list(hydrated.get("matching_image_paths") or []) if path ] return hydrated def ingest_v2_store_prepared_items_batch( connection: sqlite3.Connection, *, run_id: str, worker_id: str, entries: list[dict[str, object]], ) -> dict[str, object]: if not entries: return { "stored": 0, "store_total_ms_values": [], "serialize_ms_values": [], "spill_write_ms_values": [], "prepared_write_ms_values": [], "payload_bytes": 0, "spilled_items": 0, "spilled_bytes": 0, } store_started = time.perf_counter() active_paths = active_workspace_paths() serialized_entries: list[dict[str, object]] = [] serialize_ms_values: list[float] = [] payload_bytes_total = 0 for entry in entries: prepared_item = dict(entry["prepared_item"]) # type: ignore[index] source_fingerprint = dict(entry["source_fingerprint"]) # type: ignore[index] payload_kind = str(entry["payload_kind"]) serialize_started = time.perf_counter() spill_rel_path = None spill_path = None spill_payload = None payload_wrapper = {"prepared_item": prepared_item} if payload_kind == "mbox_message" and active_paths is not None: serialized_spill_payload = serialize_prepared_ingest_item(payload_wrapper) if ingest_v2_should_spill_prepared_payload( payload_kind=payload_kind, serialized_payload_bytes=len(serialized_spill_payload), ): spill_rel_path_obj, spill_path = ingest_v2_prepared_payload_spill_path( active_paths, run_id=run_id, work_item_id=int(entry["work_item_id"]), ) spill_rel_path = spill_rel_path_obj.as_posix() spill_payload = serialized_spill_payload payload_json = compact_json_text({"prepared_item_spilled": True}) payload_bytes = len(serialized_spill_payload) else: payload_json = compact_json_text(ingest_v2_json_safe_value(payload_wrapper)) payload_bytes = len(payload_json.encode("utf-8")) else: payload_json = compact_json_text(ingest_v2_json_safe_value(payload_wrapper)) payload_bytes = len(payload_json.encode("utf-8")) source_fingerprint_json = compact_json_text(ingest_v2_json_safe_value(source_fingerprint)) serialize_ms_values.append(ingest_v2_elapsed_ms(serialize_started)) payload_bytes_total += payload_bytes prepare_error = prepared_item.get("prepare_error") serialized_entries.append( { "work_item_id": int(entry["work_item_id"]), "payload_kind": payload_kind, "prepared_item": prepared_item, "payload_json": payload_json, "spill_rel_path": spill_rel_path, "spill_path": spill_path, "spill_payload": spill_payload, "payload_bytes": payload_bytes, "source_fingerprint_json": source_fingerprint_json, "error_json": compact_json_text({"prepare_error": prepare_error} if prepare_error else {}), "prepare_error": prepare_error, } ) now = utc_now() spill_write_ms_values: list[float] = [] written_spill_paths: list[Path] = [] prepared_write_ms = 0.0 spilled_items = 0 spilled_bytes = 0 connection.execute("BEGIN") try: run_row = require_ingest_v2_run_row(connection, run_id) if str(run_row["status"]) in INGEST_V2_TERMINAL_STATUSES or run_row["cancel_requested_at"] is not None: connection.rollback() return { "stored": 0, "store_total_ms_values": [], "serialize_ms_values": serialize_ms_values, "spill_write_ms_values": [], "prepared_write_ms_values": [], "payload_bytes": payload_bytes_total, "spilled_items": 0, "spilled_bytes": 0, } work_item_ids = [int(entry["work_item_id"]) for entry in serialized_entries] placeholders = ",".join("?" for _ in work_item_ids) item_rows = connection.execute( f""" SELECT id, status, lease_owner FROM ingest_work_items WHERE run_id = ? AND id IN ({placeholders}) """, (run_id, *work_item_ids), ).fetchall() eligible_ids = { int(row["id"]) for row in item_rows if row["status"] == "leased" and row["lease_owner"] == worker_id } if not eligible_ids: connection.rollback() return { "stored": 0, "store_total_ms_values": [], "serialize_ms_values": serialize_ms_values, "spill_write_ms_values": [], "prepared_write_ms_values": [], "payload_bytes": payload_bytes_total, "spilled_items": 0, "spilled_bytes": 0, } stored = 0 stored_entries = [ entry for entry in serialized_entries if int(entry["work_item_id"]) in eligible_ids ] for entry in stored_entries: spill_path = entry.get("spill_path") spill_payload = entry.get("spill_payload") if isinstance(spill_path, Path) and isinstance(spill_payload, bytes): spill_write_started = time.perf_counter() spill_path.parent.mkdir(parents=True, exist_ok=True) spill_path.write_bytes(spill_payload) spill_write_ms_values.append(ingest_v2_elapsed_ms(spill_write_started)) written_spill_paths.append(spill_path) spilled_items += 1 spilled_bytes += int(entry["payload_bytes"] or 0) db_write_started = time.perf_counter() payload_kind_counts: dict[str, int] = {} spilled_payload_kind_counts: dict[str, int] = {} for entry in stored_entries: work_item_id = int(entry["work_item_id"]) prepared_item = dict(entry["prepared_item"]) # type: ignore[index] prepare_error = entry["prepare_error"] payload_kind = str(entry["payload_kind"]) payload_kind_counts[payload_kind] = int(payload_kind_counts.get(payload_kind) or 0) + 1 if entry.get("spill_rel_path"): spilled_payload_kind_counts[payload_kind] = int(spilled_payload_kind_counts.get(payload_kind) or 0) + 1 connection.execute( """ INSERT INTO ingest_prepared_items ( run_id, work_item_id, payload_kind, payload_json, spill_rel_path, payload_bytes, source_fingerprint_json, prepared_at, error_json ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(work_item_id) DO UPDATE SET payload_kind = excluded.payload_kind, payload_json = excluded.payload_json, spill_rel_path = excluded.spill_rel_path, payload_bytes = excluded.payload_bytes, source_fingerprint_json = excluded.source_fingerprint_json, prepared_at = excluded.prepared_at, error_json = excluded.error_json """, ( run_id, work_item_id, payload_kind, entry["payload_json"], entry["spill_rel_path"], entry["payload_bytes"], entry["source_fingerprint_json"], now, entry["error_json"], ), ) connection.execute( """ UPDATE ingest_work_items SET status = 'prepared', lease_owner = NULL, lease_expires_at = NULL, updated_at = ?, last_error = ? WHERE run_id = ? AND id = ? """, (now, str(prepare_error) if prepare_error else None, run_id, work_item_id), ) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=worker_id, event_type="prepare_item", work_item_id=work_item_id, phase="prepare", duration_ms=prepared_item.get("prepare_total_step_ms", prepared_item.get("prepare_ms")), details={ "payload_kind": payload_kind, "rel_path": prepared_item.get("rel_path"), "prepare_error": prepare_error, }, created_at=now, ) if payload_kind == "mbox_message": ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=worker_id, event_type="prepare_mbox_message", work_item_id=work_item_id, phase="prepare", duration_ms=prepared_item.get("prepare_total_step_ms", prepared_item.get("prepare_ms")), details={ "source_rel_path": str(prepared_item.get("source_rel_path") or ""), "source_item_id": str(prepared_item.get("source_item_id") or ""), "rel_path": prepared_item.get("rel_path"), "skip": bool(prepared_item.get("skip")), "attachment_count": len(list(prepared_item.get("attachments") or [])), "prepare_hash_ms": float(prepared_item.get("prepare_hash_ms") or 0.0), "prepare_extract_ms": float(prepared_item.get("prepare_extract_ms") or 0.0), "prepare_chunk_ms": float(prepared_item.get("prepare_chunk_ms") or 0.0), "prepare_total_step_ms": float(prepared_item.get("prepare_total_step_ms") or 0.0), "prepare_error": prepare_error, }, created_at=now, ) elif payload_kind == "pst_message": attachment_count, attachment_bytes = ingest_v2_attachment_stats(prepared_item.get("attachments")) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=worker_id, event_type="prepare_pst_message", work_item_id=work_item_id, phase="prepare", duration_ms=prepared_item.get("prepare_total_step_ms", prepared_item.get("prepare_ms")), details={ "source_rel_path": str(prepared_item.get("source_rel_path") or ""), "source_item_id": str(prepared_item.get("source_item_id") or ""), "rel_path": prepared_item.get("rel_path"), "skip": bool(prepared_item.get("skip")), "attachment_count": attachment_count, "attachment_bytes": attachment_bytes, "prepare_hash_ms": float(prepared_item.get("prepare_hash_ms") or 0.0), "prepare_extract_ms": float(prepared_item.get("prepare_extract_ms") or 0.0), "prepare_chunk_ms": float(prepared_item.get("prepare_chunk_ms") or 0.0), "prepare_total_step_ms": float(prepared_item.get("prepare_total_step_ms") or 0.0), "prepare_error": prepare_error, }, created_at=now, ) stored += 1 prepared_write_ms = ingest_v2_elapsed_ms(db_write_started) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=worker_id, event_type="prepare_store_batch", phase="prepare", duration_ms=ingest_v2_elapsed_ms(store_started), details={ "stored": stored, "payload_bytes": payload_bytes_total, "serialize_ms_total": round(sum(float(value) for value in serialize_ms_values), 3), "spill_write_ms_total": round(sum(float(value) for value in spill_write_ms_values), 3), "prepared_write_ms_total": round(float(prepared_write_ms), 3), "spilled_items": spilled_items, "spilled_bytes": spilled_bytes, "payload_kind_counts": dict(sorted(payload_kind_counts.items())), "spilled_payload_kind_counts": dict(sorted(spilled_payload_kind_counts.items())), }, created_at=now, ) connection.execute( """ UPDATE ingest_runs SET last_heartbeat_at = ? WHERE run_id = ? """, (now, run_id), ) connection.commit() return { "stored": stored, "store_total_ms_values": [ingest_v2_elapsed_ms(store_started)], "serialize_ms_values": serialize_ms_values, "spill_write_ms_values": spill_write_ms_values, "prepared_write_ms_values": [prepared_write_ms], "payload_bytes": payload_bytes_total, "spilled_items": spilled_items, "spilled_bytes": spilled_bytes, } except Exception: connection.rollback() for spill_path in written_spill_paths: remove_file_if_exists(spill_path) raise def ingest_v2_store_prepared_item( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, worker_id: str, payload_kind: str, prepared_item: dict[str, object], source_fingerprint: dict[str, object], ) -> bool: result = ingest_v2_store_prepared_items_batch( connection, run_id=run_id, worker_id=worker_id, entries=[ { "work_item_id": work_item_id, "payload_kind": payload_kind, "prepared_item": prepared_item, "source_fingerprint": source_fingerprint, } ], ) return int(result.get("stored") or 0) > 0 def ingest_v2_maybe_advance_after_prepare(connection: sqlite3.Connection, *, run_id: str) -> bool: row = require_ingest_v2_run_row(connection, run_id) if str(row["phase"]) != "preparing" or row["cancel_requested_at"] is not None: return False remaining_prepare_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND status IN ('pending', 'leased') """, (run_id,), ).fetchone()[0] or 0 ) prepared_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND status = 'prepared' """, (run_id,), ).fetchone()[0] or 0 ) if not prepared_items and remaining_prepare_items: return False if prepared_items and remaining_prepare_items and prepared_items < INGEST_V2_PREPARED_COMMIT_BATCH_TARGET: return False next_phase = "committing" if prepared_items else "finalizing" now = utc_now() connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE ingest_runs SET phase = ?, status = ?, last_heartbeat_at = ? WHERE run_id = ? AND phase = 'preparing' AND cancel_requested_at IS NULL """, (next_phase, next_phase, now, run_id), ) advanced = int(cursor.rowcount or 0) > 0 connection.commit() return advanced except Exception: connection.rollback() raise def ingest_v2_acquire_writer_lease( connection: sqlite3.Connection, *, run_id: str, writer_id: str, ) -> bool: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) stale_cutoff = ingest_v2_stale_lease_cutoff(now_dt) lease_expires_at = lease_expiration_after(INGEST_V2_WORK_ITEM_LEASE_SECONDS, now=now_dt) connection.execute("BEGIN IMMEDIATE") try: cursor = connection.execute( """ UPDATE ingest_runs SET committer_lease_owner = ?, committer_lease_expires_at = ?, committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND phase = 'committing' AND status = 'committing' AND cancel_requested_at IS NULL AND ( committer_lease_owner IS NULL OR committer_lease_expires_at IS NULL OR committer_lease_expires_at <= ? OR committer_heartbeat_at IS NULL OR committer_heartbeat_at <= ? OR committer_lease_owner = ? ) """, (writer_id, lease_expires_at, now, now, run_id, now, stale_cutoff, writer_id), ) acquired = int(cursor.rowcount or 0) > 0 connection.commit() return acquired except Exception: connection.rollback() raise def ingest_v2_release_writer_lease( connection: sqlite3.Connection, *, run_id: str, writer_id: str, ) -> None: now = utc_now() connection.execute("BEGIN") try: connection.execute( """ UPDATE ingest_runs SET committer_lease_owner = NULL, committer_lease_expires_at = NULL, committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) connection.commit() except Exception: connection.rollback() raise def ingest_v2_reclaim_stale_commit_items(connection: sqlite3.Connection, *, run_id: str) -> int: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) stale_cutoff = ingest_v2_stale_lease_cutoff(now_dt) connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'prepared', lease_owner = NULL, lease_expires_at = NULL, updated_at = ? WHERE run_id = ? AND status = 'committing' AND lease_expires_at IS NOT NULL AND ( lease_expires_at <= ? OR updated_at <= ? ) """, (now, run_id, now, stale_cutoff), ) reclaimed = int(cursor.rowcount or 0) if reclaimed: connection.execute( """ INSERT INTO ingest_worker_events ( run_id, worker_id, event_type, phase, details_json, created_at ) VALUES (?, ?, ?, ?, ?, ?) """, ( run_id, None, "reclaim_stale_commit_items", "commit", compact_json_text({"reclaimed": reclaimed}), now, ), ) connection.commit() return reclaimed except Exception: connection.rollback() raise def ingest_v2_reclaim_stale_work_items(connection: sqlite3.Connection, *, run_id: str) -> dict[str, int]: return { "prepare": ingest_v2_reclaim_stale_prepare_items(connection, run_id=run_id), "commit": ingest_v2_reclaim_stale_commit_items(connection, run_id=run_id), } def ingest_v2_load_commit_cursor(connection: sqlite3.Connection, *, run_id: str) -> dict[str, object]: cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'committing' AND cursor_key = 'loose_file_commit' """, (run_id,), ).fetchone() if cursor_row is not None: cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if isinstance(cursor, dict): cursor.setdefault("schema_version", 1) cursor.setdefault("current_ingestion_batch", None) cursor.setdefault("actions", {}) cursor.setdefault("freshness_fallbacks", 0) cursor.setdefault("production_stats", {}) cursor.setdefault("slack_stats", {}) cursor.setdefault("mbox_stats", {}) cursor.setdefault("pst_stats", {}) cursor.setdefault("container_current_ingestion_batches", {}) return cursor return { "schema_version": 1, "current_ingestion_batch": None, "actions": {}, "freshness_fallbacks": 0, "production_stats": {}, "slack_stats": {}, "mbox_stats": {}, "pst_stats": {}, "container_current_ingestion_batches": {}, } def ingest_v2_run_loose_rel_paths(connection: sqlite3.Connection, *, run_id: str) -> set[str]: rows = connection.execute( """ SELECT rel_path FROM ingest_work_items WHERE run_id = ? AND unit_type = 'loose_file' AND rel_path IS NOT NULL """, (run_id,), ).fetchall() rel_paths = {str(row["rel_path"]) for row in rows} cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'planning' AND cursor_key = 'loose_file_scan' """, (run_id,), ).fetchone() if cursor_row is not None: cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if isinstance(cursor, dict): rel_paths.update( normalize_whitespace(str(rel_path or "")).replace("\\", "/").strip("/") for rel_path in list(cursor.get("skipped_unchanged_loose_file_rel_paths") or []) if normalize_whitespace(str(rel_path or "")) ) return rel_paths def ingest_v2_run_slack_seen_rel_paths_by_root( connection: sqlite3.Connection, *, run_id: str, ) -> dict[str, set[str]]: rows = connection.execute( """ SELECT payload_json FROM ingest_work_items WHERE run_id = ? AND unit_type = 'slack_conversation' """, (run_id,), ).fetchall() seen_by_root: dict[str, set[str]] = {} for row in rows: payload = decode_json_text(row["payload_json"], default={}) or {} if not isinstance(payload, dict): continue conversation_identity = payload.get("conversation_identity") identity_root = "" if isinstance(conversation_identity, list) and len(conversation_identity) > 1: identity_root = str(conversation_identity[1]) rel_root = normalize_whitespace(str(payload.get("source_locator") or identity_root)) if not rel_root: continue rel_paths = { normalize_whitespace(str(rel_path or "")).replace("\\", "/").strip("/") for rel_path in list(payload.get("rel_paths") or []) if normalize_whitespace(str(rel_path or "")) } seen_by_root.setdefault(rel_root, set()).update(rel_paths) return seen_by_root def ingest_v2_mark_missing_slack_documents(connection: sqlite3.Connection, *, run_id: str) -> int: missing = 0 for rel_root, seen_rel_paths in sorted(ingest_v2_run_slack_seen_rel_paths_by_root(connection, run_id=run_id).items()): dataset_row = get_dataset_row( connection, source_kind=SLACK_EXPORT_SOURCE_KIND, dataset_locator=rel_root, ) if dataset_row is None: continue missing += mark_missing_slack_export_documents( connection, dataset_id=int(dataset_row["id"]), seen_rel_paths=seen_rel_paths, ) return missing def ingest_v2_run_mbox_source_rel_paths(connection: sqlite3.Connection, *, run_id: str) -> set[str]: cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'planning' AND cursor_key = 'loose_file_scan' """, (run_id,), ).fetchone() source_rel_paths: set[str] = set() if cursor_row is not None: cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if isinstance(cursor, dict): source_rel_paths.update( str(rel_path) for rel_path in list(cursor.get("scanned_mbox_source_rel_paths") or []) if normalize_whitespace(str(rel_path or "")) ) rows = connection.execute( """ SELECT rel_path FROM ingest_work_items WHERE run_id = ? AND unit_type = 'mbox_source_finalizer' AND rel_path IS NOT NULL """, (run_id,), ).fetchall() source_rel_paths.update(str(row["rel_path"]) for row in rows) return source_rel_paths def ingest_v2_redelete_retired_mbox_documents( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, ) -> int: source_rel_paths = sorted(ingest_v2_run_mbox_source_rel_paths(connection, run_id=run_id)) if not source_rel_paths: return 0 placeholders = ", ".join("?" for _ in source_rel_paths) root_occurrence_rows = connection.execute( f""" SELECT id FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND source_rel_path IN ({placeholders}) AND lifecycle_status IN ('deleted', 'missing') ORDER BY id ASC """, [MBOX_SOURCE_KIND, *source_rel_paths], ).fetchall() root_occurrence_ids = [int(row["id"]) for row in root_occurrence_rows] if not root_occurrence_ids: return 0 now = utc_now() occurrence_placeholders = ", ".join("?" for _ in root_occurrence_ids) connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'deleted', updated_at = ? WHERE lifecycle_status != 'deleted' AND (id IN ({occurrence_placeholders}) OR parent_occurrence_id IN ({occurrence_placeholders})) """, [now, *root_occurrence_ids, *root_occurrence_ids], ) document_ids = container_document_ids_for_root_occurrence_ids(connection, root_occurrence_ids) deleted = delete_documents_with_only_deleted_occurrences( connection, paths, document_ids, deleted_at=now, ) return deleted + force_delete_documents_with_no_active_occurrences( connection, document_ids, deleted_at=now, ) def force_delete_documents_with_no_active_occurrences( connection: sqlite3.Connection, document_ids: set[int], *, deleted_at: str, ) -> int: forced = 0 for document_id in sorted(int(value) for value in document_ids): active_row = connection.execute( """ SELECT 1 FROM document_occurrences WHERE document_id = ? AND lifecycle_status != 'deleted' LIMIT 1 """, (document_id,), ).fetchone() if active_row is not None: continue cursor = connection.execute( """ UPDATE documents SET lifecycle_status = 'deleted', updated_at = ? WHERE id = ? AND lifecycle_status != 'deleted' """, (deleted_at, document_id), ) forced += int(cursor.rowcount or 0) return forced def ingest_v2_run_pst_source_rel_paths(connection: sqlite3.Connection, *, run_id: str) -> set[str]: cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'planning' AND cursor_key = 'loose_file_scan' """, (run_id,), ).fetchone() source_rel_paths: set[str] = set() if cursor_row is not None: cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if isinstance(cursor, dict): source_rel_paths.update( str(rel_path) for rel_path in list(cursor.get("scanned_pst_source_rel_paths") or []) if normalize_whitespace(str(rel_path or "")) ) rows = connection.execute( """ SELECT rel_path FROM ingest_work_items WHERE run_id = ? AND unit_type = 'pst_source_finalizer' AND rel_path IS NOT NULL """, (run_id,), ).fetchall() source_rel_paths.update(str(row["rel_path"]) for row in rows) return source_rel_paths def ingest_v2_redelete_retired_pst_documents( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, ) -> int: source_rel_paths = sorted(ingest_v2_run_pst_source_rel_paths(connection, run_id=run_id)) if not source_rel_paths: return 0 placeholders = ", ".join("?" for _ in source_rel_paths) root_occurrence_rows = connection.execute( f""" SELECT id FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND source_rel_path IN ({placeholders}) AND lifecycle_status IN ('deleted', 'missing') ORDER BY id ASC """, [PST_SOURCE_KIND, *source_rel_paths], ).fetchall() root_occurrence_ids = [int(row["id"]) for row in root_occurrence_rows] if not root_occurrence_ids: return 0 now = utc_now() occurrence_placeholders = ", ".join("?" for _ in root_occurrence_ids) connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'deleted', updated_at = ? WHERE lifecycle_status != 'deleted' AND (id IN ({occurrence_placeholders}) OR parent_occurrence_id IN ({occurrence_placeholders})) """, [now, *root_occurrence_ids, *root_occurrence_ids], ) document_ids = container_document_ids_for_root_occurrence_ids(connection, root_occurrence_ids) deleted = delete_documents_with_only_deleted_occurrences( connection, paths, document_ids, deleted_at=now, ) return deleted + force_delete_documents_with_no_active_occurrences( connection, document_ids, deleted_at=now, ) def ingest_v2_load_loose_file_commit_state( connection: sqlite3.Connection, *, root: Path, run_row: sqlite3.Row, ) -> tuple[dict[str, sqlite3.Row], dict[str, list[sqlite3.Row]]]: run_id = str(run_row["run_id"]) scanned_rel_paths = ingest_v2_run_loose_rel_paths(connection, run_id=run_id) scan_scope = ingest_v2_scan_scope_from_run(root, run_row) existing_by_rel, unseen_existing_by_hash = load_loose_file_commit_state( connection, scanned_rel_paths, set(), set(), scan_scope=scan_scope, ) consumed_rows = connection.execute( """ SELECT source_occurrence_id FROM ingest_rename_consumptions WHERE run_id = ? AND source_occurrence_id IS NOT NULL """, (run_id,), ).fetchall() consumed_occurrence_ids = {int(row["source_occurrence_id"]) for row in consumed_rows} if consumed_occurrence_ids: for file_hash, rows in list(unseen_existing_by_hash.items()): filtered_rows = [row for row in rows if int(row["id"]) not in consumed_occurrence_ids] if filtered_rows: unseen_existing_by_hash[file_hash] = filtered_rows else: unseen_existing_by_hash.pop(file_hash, None) return existing_by_rel, unseen_existing_by_hash def ingest_v2_claim_next_commit_item( connection: sqlite3.Connection, *, run_id: str, writer_id: str, ) -> sqlite3.Row | None: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) lease_expires_at = lease_expiration_after(INGEST_V2_WORK_ITEM_LEASE_SECONDS, now=now_dt) connection.execute("BEGIN IMMEDIATE") try: run_row = require_ingest_v2_run_row(connection, run_id) if ( str(run_row["phase"]) != "committing" or str(run_row["status"]) != "committing" or run_row["cancel_requested_at"] is not None or run_row["committer_lease_owner"] != writer_id or not lease_is_active(run_row["committer_lease_expires_at"], now=now_dt) ): connection.rollback() return None item_row = connection.execute( """ SELECT id FROM ingest_work_items WHERE run_id = ? AND status = 'prepared' ORDER BY commit_order ASC, id ASC LIMIT 1 """, (run_id,), ).fetchone() if item_row is None: connection.commit() return None work_item_id = int(item_row["id"]) connection.execute( """ UPDATE ingest_work_items SET status = 'committing', lease_owner = ?, lease_expires_at = ?, updated_at = ? WHERE run_id = ? AND id = ? AND status = 'prepared' """, (writer_id, lease_expires_at, now, run_id, work_item_id), ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) connection.commit() except Exception: connection.rollback() raise return connection.execute( """ SELECT wi.*, pi.payload_kind, pi.payload_bytes, pi.payload_json AS prepared_payload_json, pi.spill_rel_path AS prepared_spill_rel_path, pi.source_fingerprint_json, pi.error_json FROM ingest_work_items wi JOIN ingest_prepared_items pi ON pi.work_item_id = wi.id WHERE wi.run_id = ? AND wi.id = ? AND wi.status = 'committing' AND wi.lease_owner = ? """, (run_id, work_item_id, writer_id), ).fetchone() def ingest_v2_claim_additional_slack_commit_items( connection: sqlite3.Connection, *, run_id: str, writer_id: str, source_locator: str, current_payload_bytes: int, max_items: int = INGEST_V2_SLACK_COMMIT_BATCH_MAX_ITEMS, max_payload_bytes: int = INGEST_V2_SLACK_COMMIT_BATCH_MAX_PAYLOAD_BYTES, ) -> list[sqlite3.Row]: remaining_item_slots = max(0, int(max_items or 0) - 1) if remaining_item_slots <= 0: return [] now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) lease_expires_at = lease_expiration_after(INGEST_V2_WORK_ITEM_LEASE_SECONDS, now=now_dt) connection.execute("BEGIN IMMEDIATE") try: run_row = require_ingest_v2_run_row(connection, run_id) if ( str(run_row["phase"]) != "committing" or str(run_row["status"]) != "committing" or run_row["cancel_requested_at"] is not None or run_row["committer_lease_owner"] != writer_id or not lease_is_active(run_row["committer_lease_expires_at"], now=now_dt) ): connection.rollback() return [] candidate_rows = connection.execute( """ SELECT wi.*, pi.payload_kind, pi.payload_bytes, pi.payload_json AS prepared_payload_json, pi.spill_rel_path AS prepared_spill_rel_path, pi.source_fingerprint_json, pi.error_json FROM ingest_work_items wi JOIN ingest_prepared_items pi ON pi.work_item_id = wi.id WHERE wi.run_id = ? AND wi.status = 'prepared' ORDER BY wi.commit_order ASC, wi.id ASC LIMIT ? """, (run_id, remaining_item_slots), ).fetchall() selected_rows: list[sqlite3.Row] = [] selected_ids: list[int] = [] payload_bytes_total = max(0, int(current_payload_bytes or 0)) for candidate_row in candidate_rows: if str(candidate_row["payload_kind"] or candidate_row["unit_type"] or "") != "slack_document": break try: candidate_prepared_item = ingest_v2_prepared_item_from_row(candidate_row) except Exception: break if str(candidate_prepared_item.get("source_locator") or "") != source_locator: break candidate_payload_bytes = int(candidate_row["payload_bytes"] or 0) if selected_rows and payload_bytes_total + candidate_payload_bytes > int(max_payload_bytes or 0): break if not selected_rows and payload_bytes_total + candidate_payload_bytes > int(max_payload_bytes or 0): break selected_rows.append(candidate_row) selected_ids.append(int(candidate_row["id"])) payload_bytes_total += candidate_payload_bytes if not selected_ids: connection.commit() return [] placeholders = ",".join("?" for _ in selected_ids) update_cursor = connection.execute( f""" UPDATE ingest_work_items SET status = 'committing', lease_owner = ?, lease_expires_at = ?, updated_at = ? WHERE run_id = ? AND id IN ({placeholders}) AND status = 'prepared' """, (writer_id, lease_expires_at, now, run_id, *selected_ids), ) if int(update_cursor.rowcount or 0) != len(selected_ids): connection.rollback() return [] connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) connection.commit() return selected_rows except Exception: connection.rollback() raise def ingest_v2_claim_additional_container_commit_items( connection: sqlite3.Connection, *, run_id: str, writer_id: str, source_rel_path: str, current_payload_bytes: int, source_kind: str, ) -> list[sqlite3.Row]: container_config = ingest_v2_container_source_config(source_kind) max_items = int(container_config["commit_batch_max_items"]) remaining_item_slots = max(0, max_items - 1) if remaining_item_slots <= 0: return [] now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) lease_expires_at = lease_expiration_after(INGEST_V2_WORK_ITEM_LEASE_SECONDS, now=now_dt) connection.execute("BEGIN IMMEDIATE") try: run_row = require_ingest_v2_run_row(connection, run_id) if ( str(run_row["phase"]) != "committing" or str(run_row["status"]) != "committing" or run_row["cancel_requested_at"] is not None or run_row["committer_lease_owner"] != writer_id or not lease_is_active(run_row["committer_lease_expires_at"], now=now_dt) ): connection.rollback() return [] candidate_rows = connection.execute( """ SELECT wi.*, pi.payload_kind, pi.payload_bytes, pi.payload_json AS prepared_payload_json, pi.spill_rel_path AS prepared_spill_rel_path, pi.source_fingerprint_json, pi.error_json FROM ingest_work_items wi JOIN ingest_prepared_items pi ON pi.work_item_id = wi.id WHERE wi.run_id = ? AND wi.status = 'prepared' ORDER BY wi.commit_order ASC, wi.id ASC LIMIT ? """, (run_id, remaining_item_slots), ).fetchall() selected_rows: list[sqlite3.Row] = [] selected_ids: list[int] = [] payload_kind = str(container_config["message_payload_kind"]) max_payload_bytes = int(container_config["commit_batch_max_payload_bytes"]) payload_bytes_total = max(0, int(current_payload_bytes or 0)) for candidate_row in candidate_rows: if str(candidate_row["payload_kind"] or candidate_row["unit_type"] or "") != payload_kind: break try: candidate_prepared_item = ingest_v2_prepared_item_from_row(candidate_row) except Exception: break if bool(candidate_prepared_item.get("skip")): break if str(candidate_prepared_item.get("source_rel_path") or "") != source_rel_path: break candidate_payload_bytes = int(candidate_row["payload_bytes"] or 0) if payload_bytes_total + candidate_payload_bytes > max_payload_bytes: break selected_rows.append(candidate_row) selected_ids.append(int(candidate_row["id"])) payload_bytes_total += candidate_payload_bytes if not selected_ids: connection.commit() return [] placeholders = ",".join("?" for _ in selected_ids) update_cursor = connection.execute( f""" UPDATE ingest_work_items SET status = 'committing', lease_owner = ?, lease_expires_at = ?, updated_at = ? WHERE run_id = ? AND id IN ({placeholders}) AND status = 'prepared' """, (writer_id, lease_expires_at, now, run_id, *selected_ids), ) if int(update_cursor.rowcount or 0) != len(selected_ids): connection.rollback() return [] connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) connection.commit() return selected_rows except Exception: connection.rollback() raise def ingest_v2_claim_additional_mbox_commit_items( connection: sqlite3.Connection, *, run_id: str, writer_id: str, source_rel_path: str, current_payload_bytes: int, ) -> list[sqlite3.Row]: return ingest_v2_claim_additional_container_commit_items( connection, run_id=run_id, writer_id=writer_id, source_rel_path=source_rel_path, current_payload_bytes=current_payload_bytes, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_claim_additional_pst_commit_items( connection: sqlite3.Connection, *, run_id: str, writer_id: str, source_rel_path: str, current_payload_bytes: int, ) -> list[sqlite3.Row]: return ingest_v2_claim_additional_container_commit_items( connection, run_id=run_id, writer_id=writer_id, source_rel_path=source_rel_path, current_payload_bytes=current_payload_bytes, source_kind=PST_SOURCE_KIND, ) def ingest_v2_restore_claimed_commit_items( connection: sqlite3.Connection, *, run_id: str, work_item_ids: list[int], writer_id: str, ) -> int: normalized_ids = [int(work_item_id) for work_item_id in work_item_ids if work_item_id is not None] if not normalized_ids: return 0 now = utc_now() placeholders = ",".join("?" for _ in normalized_ids) connection.execute("BEGIN") try: cursor = connection.execute( f""" UPDATE ingest_work_items SET status = 'prepared', lease_owner = NULL, lease_expires_at = NULL, updated_at = ? WHERE run_id = ? AND id IN ({placeholders}) AND status = 'committing' AND lease_owner = ? """, (now, run_id, *normalized_ids, writer_id), ) restored = int(cursor.rowcount or 0) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) connection.commit() return restored except Exception: connection.rollback() raise def ingest_v2_prepared_item_from_row(row: sqlite3.Row) -> dict[str, object]: keys = set(row.keys()) spill_rel_path = None if "prepared_spill_rel_path" in keys: spill_rel_path = normalize_whitespace(str(row["prepared_spill_rel_path"] or "")) or None elif "spill_rel_path" in keys: spill_rel_path = normalize_whitespace(str(row["spill_rel_path"] or "")) or None if spill_rel_path is not None: active_paths = active_workspace_paths() if active_paths is None: raise RetrieverError("Cannot hydrate prepared spill payload without an active workspace.") restored_payload = ingest_v2_load_spilled_prepared_payload(active_paths["state_dir"] / spill_rel_path) else: payload = decode_json_text(row["prepared_payload_json"], default={}) or {} restored_payload = ingest_v2_json_restore_value(payload) if isinstance(restored_payload, dict) and isinstance(restored_payload.get("prepared_item"), dict): return dict(restored_payload["prepared_item"]) raise RetrieverError(f"Prepared payload for work item {row['id']} is malformed.") def ingest_v2_delete_prepared_item( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, require_present: bool = True, ) -> bool: cursor = connection.execute( """ DELETE FROM ingest_prepared_items WHERE run_id = ? AND work_item_id = ? """, (run_id, work_item_id), ) deleted = int(cursor.rowcount or 0) > 0 if require_present and not deleted: raise RetrieverError(f"Prepared payload for V2 ingest work item {work_item_id} is missing.") return deleted def ingest_v2_delete_terminal_prepared_items(connection: sqlite3.Connection, *, run_id: str) -> int: cursor = connection.execute( """ DELETE FROM ingest_prepared_items WHERE run_id = ? AND work_item_id IN ( SELECT id FROM ingest_work_items WHERE run_id = ? AND status IN ('committed', 'failed', 'cancelled') ) """, (run_id, run_id), ) return int(cursor.rowcount or 0) def ingest_v2_best_effort_sqlite_cleanup(connection: sqlite3.Connection) -> None: try: connection.execute("VACUUM") except sqlite3.DatabaseError: return if current_journal_mode(connection) != "wal": return try: connection.execute("PRAGMA wal_checkpoint(TRUNCATE)") except sqlite3.DatabaseError: pass def ingest_v2_mark_commit_failed( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, message: str, ) -> bool: now = utc_now() connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'failed', lease_owner = NULL, lease_expires_at = NULL, updated_at = ?, last_error = ? WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, (now, message, run_id, work_item_id, writer_id), ) marked = int(cursor.rowcount or 0) > 0 if marked: connection.execute( """ INSERT INTO ingest_worker_events ( run_id, worker_id, event_type, work_item_id, phase, details_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( run_id, writer_id, "commit_failed", work_item_id, "commit", compact_json_text({"error": message}), now, ), ) ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.commit() return marked except Exception: connection.rollback() raise def ingest_v2_commit_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], ) -> None: now = utc_now() action = str(result.get("action") or "") affected_document_ids = [ int(result["document_id"]) ] if result.get("document_id") is not None else [] if action == "renamed" and result.get("source_occurrence_id") is not None: connection.execute( """ INSERT INTO ingest_rename_consumptions ( run_id, target_work_item_id, source_document_id, source_occurrence_id, file_hash, created_at ) VALUES (?, ?, ?, ?, ?, ?) """, ( run_id, work_item_id, result.get("source_document_id"), result.get("source_occurrence_id"), str(result.get("file_hash") or ""), now, ), ) cursor["current_ingestion_batch"] = result.get("current_ingestion_batch") actions = cursor.setdefault("actions", {}) if isinstance(actions, dict): actions[action] = int(actions.get(action) or 0) + 1 if bool(result.get("freshness_fallback")): cursor["freshness_fallbacks"] = int(cursor.get("freshness_fallbacks") or 0) + 1 ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending", ) update_cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, affected_document_ids_json = ?, artifact_manifest_json = ?, updated_at = ?, last_error = NULL WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, ( compact_json_text(affected_document_ids), compact_json_text( { "commit_action": action, "freshness_fallback": bool(result.get("freshness_fallback")), "document_id": result.get("document_id"), } ), now, run_id, work_item_id, writer_id, ), ) if int(update_cursor.rowcount or 0) != 1: raise RetrieverError(f"Could not mark V2 ingest work item {work_item_id} committed.") ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) def ingest_v2_commit_conversation_preview_work_item( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], prepared_item: dict[str, object], ) -> dict[str, object]: conversation_id = int(prepared_item.get("conversation_id") or 0) if conversation_id <= 0: raise RetrieverError("Conversation preview work item is missing conversation_id.") now = utc_now() connection.execute("BEGIN") try: connection.execute("SAVEPOINT ingest_v2_conversation_preview_item") try: refreshed = refresh_conversation_previews(connection, paths, [conversation_id]) connection.execute("RELEASE SAVEPOINT ingest_v2_conversation_preview_item") except Exception: connection.execute("ROLLBACK TO SAVEPOINT ingest_v2_conversation_preview_item") connection.execute("RELEASE SAVEPOINT ingest_v2_conversation_preview_item") raise actions = cursor.setdefault("actions", {}) if isinstance(actions, dict): actions["conversation_preview"] = int(actions.get("conversation_preview") or 0) + 1 stats = cursor.setdefault("conversation_preview_stats", {}) if isinstance(stats, dict): stats["refreshed"] = int(stats.get("refreshed") or 0) + int(refreshed) ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending", ) update_cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, affected_conversation_keys_json = ?, artifact_manifest_json = ?, updated_at = ?, last_error = NULL WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, ( compact_json_text([str(conversation_id)]), compact_json_text( { "commit_action": "conversation_preview", "conversation_id": conversation_id, "refreshed": int(refreshed), } ), now, run_id, work_item_id, writer_id, ), ) if int(update_cursor.rowcount or 0) != 1: raise RetrieverError(f"Could not mark V2 ingest conversation preview work item {work_item_id} committed.") ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) connection.commit() return { "action": "conversation_preview", "conversation_id": conversation_id, "refreshed": int(refreshed), } except Exception: connection.rollback() raise def ingest_v2_ensure_production_context( connection: sqlite3.Connection, root: Path, prepared_item: dict[str, object], ) -> tuple[int, int, int]: production_rel_root = str(prepared_item["production_rel_root"]) production_name = str(prepared_item["production_name"]) metadata_load_rel_path = str(prepared_item["metadata_load_rel_path"]) image_load_rel_path = normalize_whitespace(str(prepared_item.get("image_load_rel_path") or "")) or None dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=PRODUCTION_SOURCE_KIND, source_locator=production_rel_root, dataset_name=production_dataset_name(production_rel_root, production_name), ) production_id = upsert_production_row( connection, dataset_id=dataset_id, rel_root=production_rel_root, production_name=production_name, metadata_load_rel_path=metadata_load_rel_path, image_load_rel_path=image_load_rel_path, source_type=str(prepared_item["source_type"]), ) connection.commit() return dataset_id, dataset_source_id, production_id def ingest_v2_commit_production_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], ) -> None: now = utc_now() action = str(result.get("action") or "") affected_document_ids = [ int(result["document_id"]) ] if result.get("document_id") is not None else [] production_stats = cursor.setdefault("production_stats", {}) if isinstance(production_stats, dict): production_stats[action] = int(production_stats.get(action) or 0) + 1 production_stats["page_images_linked"] = ( int(production_stats.get("page_images_linked") or 0) + int(result.get("page_images_linked") or 0) ) ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending", ) update_cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, affected_document_ids_json = ?, affected_entity_ids_json = ?, artifact_manifest_json = ?, updated_at = ?, last_error = NULL WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, ( compact_json_text(affected_document_ids), compact_json_text([]), compact_json_text( { "commit_action": action, "document_id": result.get("document_id"), "production_id": result.get("production_id"), "control_number": result.get("control_number"), "page_images_linked": int(result.get("page_images_linked") or 0), } ), now, run_id, work_item_id, writer_id, ), ) if int(update_cursor.rowcount or 0) != 1: raise RetrieverError(f"Could not mark V2 ingest production work item {work_item_id} committed.") ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) def ingest_v2_commit_production_preview_batch( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], production_id: int, prepared_item: dict[str, object], ) -> dict[str, object]: control_number = str(prepared_item["control_number"]) prepare_error = normalize_whitespace(str(prepared_item.get("prepare_error") or "")) or None if prepare_error: raise RetrieverError(prepare_error) page_assets = [ dict(asset) for asset in list(prepared_item.get("page_assets") or []) if isinstance(asset, dict) ] connection.execute("BEGIN") try: document_row = connection.execute( """ SELECT id FROM documents WHERE production_id = ? AND control_number = ? AND lifecycle_status != 'deleted' ORDER BY id ASC LIMIT 1 """, (production_id, control_number), ).fetchone() if document_row is None: raise RetrieverError(f"Production document {control_number} is not committed yet.") document_id = int(document_row["id"]) preview_rows = write_preview_page_assets(paths, page_assets) rel_preview_paths = [str(row["rel_preview_path"]) for row in preview_rows] if rel_preview_paths: placeholders = ",".join("?" for _ in rel_preview_paths) connection.execute( f""" DELETE FROM document_previews WHERE document_id = ? AND rel_preview_path IN ({placeholders}) """, (document_id, *rel_preview_paths), ) insert_document_preview_rows(connection, document_id, preview_rows) result = { "action": "preview_batch", "control_number": control_number, "document_id": document_id, "production_id": production_id, "page_images_linked": 0, "preview_pages_generated": len(preview_rows), } ingest_v2_commit_production_work_item_hook( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, ) connection.commit() return result except Exception: connection.rollback() raise def ingest_v2_commit_slack_conversation_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], ) -> None: now = utc_now() action = str(result.get("action") or "") affected_document_ids = [ int(document_id) for document_id in list(result.get("affected_document_ids") or []) if document_id is not None ] source_locator = str(result.get("source_locator") or "") conversation_key = str(result.get("conversation_key") or "") current_batch = result.get("current_batch") if current_batch is not None: cursor["current_ingestion_batch"] = int(current_batch) slack_stats = cursor.setdefault("slack_stats", {}) if isinstance(slack_stats, dict): slack_stats["slack_conversations"] = int(slack_stats.get("slack_conversations") or 0) + 1 slack_stats["slack_documents_created"] = ( int(slack_stats.get("slack_documents_created") or 0) + int(result.get("new") or 0) ) slack_stats["slack_documents_updated"] = ( int(slack_stats.get("slack_documents_updated") or 0) + int(result.get("updated") or 0) ) ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending", ) update_cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, affected_document_ids_json = ?, affected_conversation_keys_json = ?, affected_entity_ids_json = ?, artifact_manifest_json = ?, updated_at = ?, last_error = NULL WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, ( compact_json_text(affected_document_ids), compact_json_text([f"{source_locator}:{conversation_key}"] if source_locator and conversation_key else []), compact_json_text([]), compact_json_text( { "commit_action": action, "source_kind": SLACK_EXPORT_SOURCE_KIND, "source_locator": source_locator, "conversation_key": conversation_key, "new": int(result.get("new") or 0), "updated": int(result.get("updated") or 0), "rel_paths": list(result.get("rel_paths") or []), } ), now, run_id, work_item_id, writer_id, ), ) if int(update_cursor.rowcount or 0) != 1: raise RetrieverError(f"Could not mark V2 ingest Slack work item {work_item_id} committed.") ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) def ingest_v2_commit_slack_document_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], ) -> None: now = utc_now() action = str(result.get("action") or "") affected_document_ids = [ int(document_id) for document_id in list(result.get("affected_document_ids") or []) if document_id is not None ] source_locator = str(result.get("source_locator") or "") conversation_key = str(result.get("conversation_key") or "") current_batch = result.get("current_batch") if current_batch is not None: cursor["current_ingestion_batch"] = int(current_batch) slack_stats = cursor.setdefault("slack_stats", {}) if isinstance(slack_stats, dict): if bool(result.get("conversation_lead_document")): slack_stats["slack_conversations"] = int(slack_stats.get("slack_conversations") or 0) + 1 slack_stats["slack_documents_created"] = ( int(slack_stats.get("slack_documents_created") or 0) + int(result.get("new") or 0) ) slack_stats["slack_documents_updated"] = ( int(slack_stats.get("slack_documents_updated") or 0) + int(result.get("updated") or 0) ) ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending", ) update_cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, affected_document_ids_json = ?, affected_conversation_keys_json = ?, affected_entity_ids_json = ?, artifact_manifest_json = ?, updated_at = ?, last_error = NULL WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, ( compact_json_text(affected_document_ids), compact_json_text([f"{source_locator}:{conversation_key}"] if source_locator and conversation_key else []), compact_json_text([]), compact_json_text( { "commit_action": action, "source_kind": SLACK_EXPORT_SOURCE_KIND, "source_locator": source_locator, "conversation_key": conversation_key, "document_kind": str(result.get("document_kind") or ""), "new": int(result.get("new") or 0), "updated": int(result.get("updated") or 0), "rel_paths": list(result.get("rel_paths") or []), } ), now, run_id, work_item_id, writer_id, ), ) if int(update_cursor.rowcount or 0) != 1: raise RetrieverError(f"Could not mark V2 ingest Slack work item {work_item_id} committed.") ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) def ingest_v2_commit_slack_document_batch( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, claimed_rows: list[sqlite3.Row], prepared_items: list[dict[str, object]], writer_id: str, cursor: dict[str, object], dataset_id: int, dataset_source_id: int, ) -> list[dict[str, object]]: if not claimed_rows or len(claimed_rows) != len(prepared_items): raise RetrieverError("Slack commit batch rows and prepared items must have the same non-zero length.") batch_cursor = ingest_v2_clone_cursor(cursor) current_batch = ( int(batch_cursor["current_ingestion_batch"]) if batch_cursor.get("current_ingestion_batch") is not None else None ) connection.execute("BEGIN") try: results: list[dict[str, object]] = [] for claimed_row, prepared_item in zip(claimed_rows, prepared_items): commit_result = commit_prepared_slack_document_in_transaction( connection, paths, prepared_item, dataset_id=dataset_id, dataset_source_id=dataset_source_id, current_batch=current_batch, ) if str(commit_result.get("status") or "") == "failed": raise RetrieverError(str(commit_result.get("error") or "Slack document commit failed.")) current_batch = ( int(commit_result["current_batch"]) if commit_result.get("current_batch") is not None else current_batch ) ingest_v2_commit_slack_document_work_item_hook( connection, run_id=run_id, work_item_id=int(claimed_row["id"]), writer_id=writer_id, cursor=batch_cursor, result=commit_result, ) results.append(commit_result) connection.commit() except Exception: connection.rollback() raise cursor.clear() cursor.update(batch_cursor) return results def ingest_v2_commit_container_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], source_kind: str, ) -> None: container_config = ingest_v2_container_source_config(source_kind) source_label = str(container_config["source_label"]) source_stats_key = str(container_config["source_stats_key"]) created_stat_key = str(container_config["created_stat_key"]) updated_stat_key = str(container_config["updated_stat_key"]) skipped_stat_key = str(container_config["skipped_stat_key"]) finalized_stat_key = str(container_config["finalized_stat_key"]) deleted_stat_key = str(container_config["deleted_stat_key"]) now = utc_now() action = str(result.get("action") or "") source_rel_path = str(result.get("source_rel_path") or "") affected_document_ids = ( [int(result["document_id"])] if result.get("document_id") is not None else [] ) current_batch = result.get("current_ingestion_batch") if source_rel_path and current_batch is not None: current_batches = cursor.setdefault("container_current_ingestion_batches", {}) if isinstance(current_batches, dict): current_batches[source_rel_path] = int(current_batch) source_stats = cursor.setdefault(source_stats_key, {}) if isinstance(source_stats, dict): if action == "new": source_stats[created_stat_key] = int(source_stats.get(created_stat_key) or 0) + 1 elif action == "updated": source_stats[updated_stat_key] = int(source_stats.get(updated_stat_key) or 0) + 1 for key in (skipped_stat_key, finalized_stat_key, deleted_stat_key): if result.get(key) is not None: source_stats[key] = int(source_stats.get(key) or 0) + int(result.get(key) or 0) if source_kind == MBOX_SOURCE_KIND and result.get("gmail_linked_drive_retired") is not None: source_stats["gmail_linked_drive_retired"] = ( int(source_stats.get("gmail_linked_drive_retired") or 0) + int(result.get("gmail_linked_drive_retired") or 0) ) ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending", ) update_cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, affected_document_ids_json = ?, affected_entity_ids_json = ?, artifact_manifest_json = ?, updated_at = ?, last_error = NULL WHERE run_id = ? AND id = ? AND status = 'committing' AND lease_owner = ? """, ( compact_json_text(affected_document_ids), compact_json_text([]), compact_json_text( { "commit_action": action, "document_id": result.get("document_id"), "source_kind": source_kind, "source_plan_kind": result.get("source_plan_kind"), "source_rel_path": source_rel_path, "source_item_id": result.get("source_item_id"), deleted_stat_key: int(result.get(deleted_stat_key) or 0), } ), now, run_id, work_item_id, writer_id, ), ) if int(update_cursor.rowcount or 0) != 1: raise RetrieverError( f"Could not mark V2 ingest {source_label} work item {work_item_id} committed." ) ingest_v2_delete_prepared_item( connection, run_id=run_id, work_item_id=work_item_id, ) connection.execute( """ UPDATE ingest_runs SET committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND committer_lease_owner = ? """, (now, now, run_id, writer_id), ) def ingest_v2_commit_container_message_batch( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, claimed_rows: list[sqlite3.Row], prepared_items: list[dict[str, object]], writer_id: str, cursor: dict[str, object], context: dict[str, object], source_kind: str, ) -> list[dict[str, object]]: container_config = ingest_v2_container_source_config(source_kind) source_label = str(container_config["source_label"]) default_source_plan_kind = str(container_config["source_plan_kind_default"]) message_commit_event_type = str(container_config["message_commit_event_type"]) if not claimed_rows or len(claimed_rows) != len(prepared_items): raise RetrieverError( f"{source_label} commit batch rows and prepared items must have the same non-zero length." ) batch_cursor = ingest_v2_clone_cursor(cursor) current_batch = ( int(context["current_ingestion_batch"]) if context.get("current_ingestion_batch") is not None else None ) connection.execute("BEGIN") try: results: list[dict[str, object]] = [] existing_entries_by_source_item = dict(context.get("existing_entries_by_source_item") or {}) for batch_index, (claimed_row, prepared_item) in enumerate( zip(claimed_rows, prepared_items), start=1, ): work_item_id = int(claimed_row["id"]) source_rel_path = str(prepared_item.get("source_rel_path") or "") existing_entry = dict( existing_entries_by_source_item.get(str(prepared_item["source_item_id"])) or {} ) commit_started = time.perf_counter() commit_result = commit_prepared_container_message_in_transaction( connection, paths, prepared_item, existing_entry.get("document_row"), existing_entry.get("occurrence_row"), current_ingestion_batch=current_batch, dataset_id=int(context["dataset_id"]), dataset_source_id=( int(context["dataset_source_id"]) if context.get("dataset_source_id") is not None else None ), source_kind=source_kind, source_rel_path=source_rel_path, file_type_override=source_kind, scan_started_at=str(prepared_item["scan_started_at"]), before_transaction_commit=( lambda commit_connection, result, claimed_work_item_id=work_item_id, rel_path=source_rel_path, prepared=prepared_item: ingest_v2_commit_container_work_item_hook( commit_connection, run_id=run_id, work_item_id=claimed_work_item_id, writer_id=writer_id, cursor=batch_cursor, result={ **result, "source_kind": source_kind, "source_plan_kind": str( prepared.get("source_plan_kind") or default_source_plan_kind ), "source_rel_path": rel_path, "source_item_id": str(prepared["source_item_id"]), }, source_kind=source_kind, ) ), ) current_batch = ( int(commit_result["current_ingestion_batch"]) if commit_result.get("current_ingestion_batch") is not None else current_batch ) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=writer_id, event_type=message_commit_event_type, work_item_id=work_item_id, phase="commit", duration_ms=ingest_v2_elapsed_ms(commit_started), details={ "action": str(commit_result.get("action") or ""), "source_rel_path": source_rel_path, "source_item_id": str(prepared_item.get("source_item_id") or ""), "source_plan_kind": str( prepared_item.get("source_plan_kind") or default_source_plan_kind ), "document_id": commit_result.get("document_id"), "batch_index": int(batch_index), "batch_size": int(len(claimed_rows)), }, ) results.append(commit_result) connection.commit() except Exception: connection.rollback() raise context["current_ingestion_batch"] = current_batch cursor.clear() cursor.update(batch_cursor) return results def ingest_v2_commit_mbox_message_batch( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, claimed_rows: list[sqlite3.Row], prepared_items: list[dict[str, object]], writer_id: str, cursor: dict[str, object], context: dict[str, object], ) -> list[dict[str, object]]: return ingest_v2_commit_container_message_batch( connection, paths, run_id=run_id, claimed_rows=claimed_rows, prepared_items=prepared_items, writer_id=writer_id, cursor=cursor, context=context, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_commit_pst_message_batch( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, claimed_rows: list[sqlite3.Row], prepared_items: list[dict[str, object]], writer_id: str, cursor: dict[str, object], context: dict[str, object], ) -> list[dict[str, object]]: return ingest_v2_commit_container_message_batch( connection, paths, run_id=run_id, claimed_rows=claimed_rows, prepared_items=prepared_items, writer_id=writer_id, cursor=cursor, context=context, source_kind=PST_SOURCE_KIND, ) def ingest_v2_commit_mbox_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], ) -> None: ingest_v2_commit_container_work_item_hook( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_commit_pst_work_item_hook( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], result: dict[str, object], ) -> None: ingest_v2_commit_container_work_item_hook( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, source_kind=PST_SOURCE_KIND, ) def ingest_v2_ensure_container_commit_context( connection: sqlite3.Connection, *, prepared_item: dict[str, object], cursor: dict[str, object], contexts: dict[str, dict[str, object]], source_kind: str, ) -> dict[str, object]: container_config = ingest_v2_container_source_config(source_kind) source_rel_path = str(prepared_item["source_rel_path"]) if source_rel_path in contexts: return contexts[source_rel_path] dataset_name_builder = container_config["dataset_name_builder"] assert callable(dataset_name_builder) dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=source_kind, source_locator=source_rel_path, dataset_name=dataset_name_builder(source_rel_path), ) if connection.in_transaction: connection.commit() connection.execute("BEGIN") try: write_container_source_scan_started( connection, dataset_id=dataset_id, source_kind=source_kind, source_rel_path=source_rel_path, file_size=( int(prepared_item["source_file_size"]) if prepared_item.get("source_file_size") is not None else None ), file_mtime=( str(prepared_item["source_file_mtime"]) if prepared_item.get("source_file_mtime") is not None else None ), file_hash=str(prepared_item.get("source_file_hash") or ""), scan_started_at=str(prepared_item["scan_started_at"]), ) connection.commit() except Exception: connection.rollback() raise current_batches = cursor.setdefault("container_current_ingestion_batches", {}) current_batch = None if isinstance(current_batches, dict) and current_batches.get(source_rel_path) is not None: current_batch = int(current_batches[source_rel_path]) context = { "dataset_id": int(dataset_id), "dataset_source_id": int(dataset_source_id) if dataset_source_id is not None else None, "current_ingestion_batch": current_batch, "existing_entries_by_source_item": existing_container_entries_by_source_item( connection, source_kind=source_kind, source_rel_path=source_rel_path, ), } contexts[source_rel_path] = context return context def ingest_v2_ensure_mbox_commit_context( connection: sqlite3.Connection, root: Path, *, prepared_item: dict[str, object], cursor: dict[str, object], contexts: dict[str, dict[str, object]], ) -> dict[str, object]: _ = root return ingest_v2_ensure_container_commit_context( connection, prepared_item=prepared_item, cursor=cursor, contexts=contexts, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_ensure_pst_commit_context( connection: sqlite3.Connection, root: Path, *, prepared_item: dict[str, object], cursor: dict[str, object], contexts: dict[str, dict[str, object]], ) -> dict[str, object]: _ = root return ingest_v2_ensure_container_commit_context( connection, prepared_item=prepared_item, cursor=cursor, contexts=contexts, source_kind=PST_SOURCE_KIND, ) def ingest_v2_container_failed_message_count( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, source_kind: str, ) -> int: payload_kind = str(ingest_v2_container_source_config(source_kind)["message_payload_kind"]) prefix = f"{source_rel_path}:" row = connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND unit_type = ? AND status = 'failed' AND SUBSTR(COALESCE(source_key, ''), 1, ?) = ? """, (run_id, payload_kind, len(prefix), prefix), ).fetchone() return int(row[0] or 0) if row is not None else 0 def ingest_v2_mbox_failed_message_count( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, ) -> int: return ingest_v2_container_failed_message_count( connection, run_id=run_id, source_rel_path=source_rel_path, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_pst_failed_message_count( connection: sqlite3.Connection, *, run_id: str, source_rel_path: str, ) -> int: return ingest_v2_container_failed_message_count( connection, run_id=run_id, source_rel_path=source_rel_path, source_kind=PST_SOURCE_KIND, ) def ingest_v2_container_work_item_kind(payload_kind: str, unit_type: str) -> tuple[str, str] | None: for source_kind in (MBOX_SOURCE_KIND, PST_SOURCE_KIND): container_config = ingest_v2_container_source_config(source_kind) message_payload_kind = str(container_config["message_payload_kind"]) source_finalizer_payload_kind = str(container_config["source_finalizer_payload_kind"]) if payload_kind == message_payload_kind or unit_type == message_payload_kind: return source_kind, "message" if payload_kind == source_finalizer_payload_kind or unit_type == source_finalizer_payload_kind: return source_kind, "finalizer" return None def ingest_v2_raise_prepare_error(prepared_item: dict[str, object]) -> None: prepare_error = normalize_whitespace(str(prepared_item.get("prepare_error") or "")) or None if prepare_error: raise RetrieverError(prepare_error) def ingest_v2_container_message_batch_function( source_kind: str, ) -> Callable[..., list[dict[str, object]]]: if source_kind == MBOX_SOURCE_KIND: return ingest_v2_commit_mbox_message_batch if source_kind == PST_SOURCE_KIND: return ingest_v2_commit_pst_message_batch raise RetrieverError(f"Unsupported container source kind: {source_kind!r}") def ingest_v2_expand_container_commit_batch( connection: sqlite3.Connection, *, run_id: str, writer_id: str, claimed_row: sqlite3.Row, prepared_item: dict[str, object], source_kind: str, prepared_decode_ms_values: list[float], ) -> tuple[list[sqlite3.Row], list[dict[str, object]]]: batch_claimed_rows = [claimed_row] batch_prepared_items = [prepared_item] additional_rows = ingest_v2_claim_additional_container_commit_items( connection, run_id=run_id, writer_id=writer_id, source_rel_path=str(prepared_item["source_rel_path"]), current_payload_bytes=int(claimed_row["payload_bytes"] or 0), source_kind=source_kind, ) if not additional_rows: return batch_claimed_rows, batch_prepared_items additional_prepared_items: list[dict[str, object]] = [] additional_decode_started = time.perf_counter() decode_failed = False for additional_row in additional_rows: try: additional_prepared_items.append(ingest_v2_prepared_item_from_row(additional_row)) except Exception: decode_failed = True break decoded_count = len(additional_prepared_items) if decoded_count: additional_decode_ms = ingest_v2_elapsed_ms(additional_decode_started) prepared_decode_ms_values.extend([additional_decode_ms / decoded_count] * decoded_count) if decode_failed: ingest_v2_restore_claimed_commit_items( connection, run_id=run_id, work_item_ids=[int(row["id"]) for row in additional_rows], writer_id=writer_id, ) return batch_claimed_rows, batch_prepared_items batch_claimed_rows.extend(additional_rows) batch_prepared_items.extend(additional_prepared_items) return batch_claimed_rows, batch_prepared_items def ingest_v2_commit_skipped_container_message( connection: sqlite3.Connection, *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], prepared_item: dict[str, object], source_kind: str, ) -> dict[str, object]: container_config = ingest_v2_container_source_config(source_kind) message_commit_event_type = str(container_config["message_commit_event_type"]) default_source_plan_kind = str(container_config["source_plan_kind_default"]) source_rel_path = str(prepared_item.get("source_rel_path") or "") source_item_id = str(prepared_item.get("source_item_id") or "") connection.execute("BEGIN") try: commit_result = { "action": "skipped", "source_kind": source_kind, "source_plan_kind": str(prepared_item.get("source_plan_kind") or default_source_plan_kind), "source_rel_path": source_rel_path, "source_item_id": source_item_id, "current_ingestion_batch": None, "document_id": None, } ingest_v2_commit_container_work_item_hook( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=commit_result, source_kind=source_kind, ) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=writer_id, event_type=message_commit_event_type, work_item_id=work_item_id, phase="commit", duration_ms=0.0, details={ "action": "skipped", "source_rel_path": source_rel_path, "source_item_id": source_item_id, "source_plan_kind": str( prepared_item.get("source_plan_kind") or default_source_plan_kind ), "document_id": None, "batch_index": 1, "batch_size": 1, }, ) connection.commit() except Exception: connection.rollback() raise return commit_result def ingest_v2_commit_container_message_work_item( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, claimed_row: sqlite3.Row, prepared_item: dict[str, object], writer_id: str, cursor: dict[str, object], container_contexts: dict[str, dict[str, dict[str, object]]], source_kind: str, prepared_decode_ms_values: list[float], ) -> dict[str, object]: if prepared_item.get("skip"): skipped_result = ingest_v2_commit_skipped_container_message( connection, run_id=run_id, work_item_id=int(claimed_row["id"]), writer_id=writer_id, cursor=cursor, prepared_item=prepared_item, source_kind=source_kind, ) return { "batch_results": [skipped_result], "batch_count": 0, "batch_items": 0, "batch_fallbacks": 0, "timing_item_count": 1, } source_contexts = container_contexts.setdefault(source_kind, {}) context = ingest_v2_ensure_container_commit_context( connection, prepared_item=prepared_item, cursor=cursor, contexts=source_contexts, source_kind=source_kind, ) batch_claimed_rows, batch_prepared_items = ingest_v2_expand_container_commit_batch( connection, run_id=run_id, writer_id=writer_id, claimed_row=claimed_row, prepared_item=prepared_item, source_kind=source_kind, prepared_decode_ms_values=prepared_decode_ms_values, ) commit_batch = ingest_v2_container_message_batch_function(source_kind) batch_count = 0 batch_items = 0 batch_fallbacks = 0 try: batch_results = commit_batch( connection, paths, run_id=run_id, claimed_rows=batch_claimed_rows, prepared_items=batch_prepared_items, writer_id=writer_id, cursor=cursor, context=context, ) if len(batch_results) > 1: batch_count = 1 batch_items = len(batch_results) except Exception: rollback_open_transaction(connection) if len(batch_claimed_rows) <= 1: raise batch_fallbacks = 1 ingest_v2_restore_claimed_commit_items( connection, run_id=run_id, work_item_ids=[int(row["id"]) for row in batch_claimed_rows[1:]], writer_id=writer_id, ) batch_results = commit_batch( connection, paths, run_id=run_id, claimed_rows=[claimed_row], prepared_items=[prepared_item], writer_id=writer_id, cursor=cursor, context=context, ) return { "batch_results": batch_results, "batch_count": batch_count, "batch_items": batch_items, "batch_fallbacks": batch_fallbacks, "timing_item_count": max(1, len(batch_results)), } def ingest_v2_commit_container_source_finalizer( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], prepared_item: dict[str, object], source_kind: str, ) -> dict[str, object]: container_config = ingest_v2_container_source_config(source_kind) source_label = str(container_config["source_label"]) default_source_plan_kind = str(container_config["source_plan_kind_default"]) skipped_stat_key = str(container_config["skipped_stat_key"]) finalized_stat_key = str(container_config["finalized_stat_key"]) deleted_stat_key = str(container_config["deleted_stat_key"]) dataset_name_builder = container_config["dataset_name_builder"] assert callable(dataset_name_builder) finalizer_started = time.perf_counter() source_rel_path = str(prepared_item["source_rel_path"]) failed_messages = ingest_v2_container_failed_message_count( connection, run_id=run_id, source_rel_path=source_rel_path, source_kind=source_kind, ) if failed_messages: raise RetrieverError( f"Cannot finalize {source_label} source {source_rel_path}: " f"{failed_messages} message work item(s) failed." ) dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=source_kind, source_locator=source_rel_path, dataset_name=dataset_name_builder(source_rel_path), ) if connection.in_transaction: connection.commit() source_file_size = ( int(prepared_item["source_file_size"]) if prepared_item.get("source_file_size") is not None else None ) source_file_mtime = ( str(prepared_item["source_file_mtime"]) if prepared_item.get("source_file_mtime") is not None else None ) source_file_hash = str(prepared_item.get("source_file_hash") or "") scan_started_at = str(prepared_item["scan_started_at"]) skip_source = bool(prepared_item.get("skip_source")) source_plan_kind = str(prepared_item.get("source_plan_kind") or default_source_plan_kind) connection.execute("BEGIN") try: linked_drive_retired = 0 linked_drive_rel_paths: set[str] = set() if source_kind == MBOX_SOURCE_KIND: linked_drive_rel_paths = { normalize_whitespace(str(rel_path or "")).replace("\\", "/").strip("/") for rel_path in list(prepared_item.get("linked_drive_rel_paths") or []) if normalize_whitespace(str(rel_path or "")) } if source_kind == MBOX_SOURCE_KIND and source_plan_kind == "gmail" and linked_drive_rel_paths: linked_drive_retired = retire_standalone_filesystem_documents_by_rel_paths( connection, paths, rel_paths=linked_drive_rel_paths, ) if skip_source: mark_container_source_documents_active( connection, source_kind=source_kind, source_rel_path=source_rel_path, seen_at=scan_started_at, ) assign_dataset_to_container_documents( connection, source_kind=source_kind, source_rel_path=source_rel_path, dataset_id=dataset_id, dataset_source_id=dataset_source_id, ) messages_deleted = 0 action = "skipped" source_stats = {skipped_stat_key: 1} scan_completed_at = scan_started_at else: messages_deleted = retire_unseen_container_messages( connection, paths, source_kind=source_kind, source_rel_path=source_rel_path, scan_started_at=scan_started_at, ) action = "finalized" source_stats = {finalized_stat_key: 1} scan_completed_at = next_monotonic_utc_timestamp([scan_started_at]) message_count = int(prepared_item.get("message_count") or 0) if source_kind == PST_SOURCE_KIND and not skip_source: message_count = len( container_root_occurrence_rows_for_source( connection, source_kind=source_kind, source_rel_path=source_rel_path, ) ) write_container_source_scan_completed( connection, dataset_id=dataset_id, source_kind=source_kind, source_rel_path=source_rel_path, file_size=source_file_size, file_mtime=source_file_mtime, file_hash=source_file_hash, message_count=message_count, scan_started_at=scan_started_at, scan_completed_at=scan_completed_at, ) result = { "action": action, "source_kind": source_kind, "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "current_ingestion_batch": None, "document_id": None, deleted_stat_key: messages_deleted, **source_stats, } if source_kind == MBOX_SOURCE_KIND: result["gmail_linked_drive_retired"] = linked_drive_retired ingest_v2_commit_container_work_item_hook( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, source_kind=source_kind, ) if source_kind == PST_SOURCE_KIND: ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=writer_id, event_type="commit_pst_source_finalizer", work_item_id=work_item_id, phase="commit", duration_ms=ingest_v2_elapsed_ms(finalizer_started), details={ "action": action, "source_rel_path": source_rel_path, "source_plan_kind": source_plan_kind, "message_count": message_count, deleted_stat_key: messages_deleted, "skip_source": skip_source, }, ) connection.commit() except Exception: connection.rollback() raise return result def ingest_v2_commit_mbox_source_finalizer( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], prepared_item: dict[str, object], ) -> dict[str, object]: return ingest_v2_commit_container_source_finalizer( connection, paths, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, prepared_item=prepared_item, source_kind=MBOX_SOURCE_KIND, ) def ingest_v2_commit_pst_source_finalizer( connection: sqlite3.Connection, paths: dict[str, Path], *, run_id: str, work_item_id: int, writer_id: str, cursor: dict[str, object], prepared_item: dict[str, object], ) -> dict[str, object]: return ingest_v2_commit_container_source_finalizer( connection, paths, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, prepared_item=prepared_item, source_kind=PST_SOURCE_KIND, ) def ingest_v2_maybe_advance_after_commit(connection: sqlite3.Connection, *, run_id: str) -> bool: row = require_ingest_v2_run_row(connection, run_id) if str(row["phase"]) != "committing" or row["cancel_requested_at"] is not None: return False remaining_commit_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND status IN ('prepared', 'committing') """, (run_id,), ).fetchone()[0] or 0 ) if remaining_commit_items: return False remaining_prepare_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND status IN ('pending', 'leased') """, (run_id,), ).fetchone()[0] or 0 ) next_phase = "preparing" if remaining_prepare_items else "finalizing" cursor = ingest_v2_load_commit_cursor(connection, run_id=run_id) now = utc_now() connection.execute("BEGIN") try: ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", cursor=cursor, status="pending" if remaining_prepare_items else "complete", ) update_cursor = connection.execute( """ UPDATE ingest_runs SET phase = ?, status = ?, committer_lease_owner = NULL, committer_lease_expires_at = NULL, committer_heartbeat_at = ?, last_heartbeat_at = ? WHERE run_id = ? AND phase = 'committing' AND cancel_requested_at IS NULL """, (next_phase, next_phase, now, now, run_id), ) advanced = int(update_cursor.rowcount or 0) > 0 connection.commit() return advanced except Exception: connection.rollback() raise def ingest_v2_planned_production_roots_by_rel_root(connection: sqlite3.Connection, *, run_id: str) -> dict[str, dict[str, object]]: cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'planning' AND cursor_key = 'loose_file_scan' """, (run_id,), ).fetchone() if cursor_row is None: return {} cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if not isinstance(cursor, dict): return {} roots_by_rel_root = dict(cursor.get("production_roots_by_rel_root") or {}) planned_roots = [str(rel_root) for rel_root in list(cursor.get("planned_production_roots") or [])] return { rel_root: dict(roots_by_rel_root.get(rel_root) or {}) for rel_root in planned_roots if roots_by_rel_root.get(rel_root) } def ingest_v2_finalize_production_root( connection: sqlite3.Connection, paths: dict[str, Path], root: Path, *, production_payload: dict[str, object], ) -> dict[str, int]: dataset_id, _dataset_source_id, production_id = ingest_v2_ensure_production_context( connection, root, { "production_rel_root": production_payload["rel_root"], "production_name": production_payload["production_name"], "metadata_load_rel_path": production_payload["metadata_load_rel_path"], "image_load_rel_path": production_payload.get("image_load_rel_path"), "source_type": production_payload["source_type"], }, ) seen_control_numbers = {str(control_number) for control_number in list(production_payload.get("seen_control_numbers") or [])} retired = 0 existing_rows = connection.execute( """ SELECT * FROM documents WHERE production_id = ? """, (production_id,), ).fetchall() for row in existing_rows: control_number = str(row["control_number"] or "") if control_number and control_number in seen_control_numbers: continue connection.execute("BEGIN") try: cleanup_document_artifacts(paths, connection, row) delete_document_related_rows(connection, int(row["id"])) connection.execute( """ UPDATE documents SET lifecycle_status = 'deleted', parent_document_id = NULL, dataset_id = ?, updated_at = ? WHERE id = ? """, (dataset_id, utc_now(), row["id"]), ) connection.commit() retired += 1 except Exception: connection.rollback() raise connection.execute("BEGIN") try: parent_link_updates = update_production_family_relationships(connection, production_id) attachment_preview_updates = 0 preview_document_rows = connection.execute( """ SELECT DISTINCT documents.id FROM documents JOIN document_previews ON document_previews.document_id = documents.id WHERE documents.production_id = ? AND documents.lifecycle_status != 'deleted' AND document_previews.preview_type = 'html' ORDER BY documents.id ASC """, (production_id,), ).fetchall() for preview_document_row in preview_document_rows: attachment_preview_updates += sync_document_attachment_preview_links( connection, paths, int(preview_document_row["id"]), ) connection.commit() except Exception: connection.rollback() raise families_reconstructed = len( connection.execute( """ SELECT id FROM documents WHERE production_id = ? AND parent_document_id IS NOT NULL AND lifecycle_status != 'deleted' """, (production_id,), ).fetchall() ) return { "retired": retired, "families_reconstructed": families_reconstructed, "parent_link_updates": int(parent_link_updates), "attachment_preview_updates": int(attachment_preview_updates), } def ingest_v2_load_finalize_cursor(connection: sqlite3.Connection, *, run_id: str) -> dict[str, object]: cursor_row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = 'finalizing' AND cursor_key = 'loose_file_finalize' """, (run_id,), ).fetchone() if cursor_row is not None: cursor = decode_json_text(cursor_row["cursor_json"], default={}) or {} if isinstance(cursor, dict): cursor.setdefault("schema_version", 1) cursor.setdefault("stage", "production") cursor.setdefault("pending_production_rel_roots", []) cursor.setdefault("production_roots_by_rel_root", {}) cursor.setdefault("production_finalized_roots", []) cursor.setdefault("production_stats", {}) cursor.setdefault("filesystem_missing", 0) cursor.setdefault("mbox_sources_missing", 0) cursor.setdefault("mbox_documents_missing", 0) cursor.setdefault("mbox_documents_redeleted", 0) cursor.setdefault("conversation_assignment", {}) cursor.setdefault("conversation_preview_active_count", 0) cursor.setdefault("conversation_preview_target_count", 0) cursor.setdefault("conversation_preview_skipped_fresh", 0) cursor.setdefault("conversation_preview_forced_count", 0) cursor.setdefault("conversation_preview_work_items_planned", 0) cursor.setdefault("conversation_previews_refreshed", 0) cursor.setdefault("pruned_unused_filesystem_dataset", False) return cursor production_roots_by_rel_root = ingest_v2_planned_production_roots_by_rel_root(connection, run_id=run_id) return { "schema_version": 1, "stage": "production" if production_roots_by_rel_root else "missing", "pending_production_rel_roots": sorted(production_roots_by_rel_root), "production_roots_by_rel_root": production_roots_by_rel_root, "production_finalized_roots": [], "production_stats": {}, "filesystem_missing": 0, "mbox_sources_missing": 0, "mbox_documents_missing": 0, "mbox_documents_redeleted": 0, "conversation_assignment": {}, "conversation_preview_active_count": 0, "conversation_preview_target_count": 0, "conversation_preview_skipped_fresh": 0, "conversation_preview_forced_count": 0, "conversation_preview_work_items_planned": 0, "conversation_previews_refreshed": 0, "pruned_unused_filesystem_dataset": False, } def ingest_v2_save_finalize_cursor( connection: sqlite3.Connection, *, run_id: str, cursor: dict[str, object], status: str, ) -> None: ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="finalizing", cursor_key="loose_file_finalize", cursor=cursor, status=status, ) connection.execute( """ UPDATE ingest_runs SET last_heartbeat_at = ? WHERE run_id = ? """, (utc_now(), run_id), ) def active_ingest_v2_run_row(connection: sqlite3.Connection) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM ingest_runs WHERE status NOT IN ('completed', 'canceled', 'failed') ORDER BY created_at DESC, id DESC LIMIT 1 """ ).fetchone() def require_ingest_v2_run_row(connection: sqlite3.Connection, run_id: str) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM ingest_runs WHERE run_id = ? """, (run_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown resumable ingest run: {run_id}") return row def latest_ingest_v2_run_row(connection: sqlite3.Connection) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM ingest_runs ORDER BY created_at DESC, id DESC LIMIT 1 """ ).fetchone() def ingest_v2_conflict_payload(root: Path, active_row: sqlite3.Row, *, message: str) -> dict[str, object]: run_id = str(active_row["run_id"]) quoted_root = shlex.quote(str(root)) quoted_run_id = shlex.quote(run_id) return { "ok": False, "error": "active_ingest_run", "active_run_id": run_id, "message": message, "status_command": f"ingest-status {quoted_root} --run-id {quoted_run_id}", "cancel_command": f"ingest-cancel {quoted_root} --run-id {quoted_run_id}", } def raise_if_ingest_v2_active(connection: sqlite3.Connection, root: Path, *, command_name: str) -> None: active_row = active_ingest_v2_run_row(connection) if active_row is None: return raise RetrieverStructuredError( f"{command_name} cannot run while resumable ingest run {active_row['run_id']} is active.", ingest_v2_conflict_payload( root, active_row, message=f"{command_name} cannot run while a resumable ingest run is active.", ), ) def ingest_v2_status_counts(connection: sqlite3.Connection, *, run_id: str) -> dict[str, int]: counts = {status: 0 for status in INGEST_V2_WORK_ITEM_STATUSES} rows = connection.execute( """ SELECT status, COUNT(*) AS count FROM ingest_work_items WHERE run_id = ? GROUP BY status """, (run_id,), ).fetchall() for row in rows: counts[str(row["status"])] = int(row["count"] or 0) return counts def ingest_v2_unit_type_counts(connection: sqlite3.Connection, *, run_id: str) -> dict[str, dict[str, int]]: rows = connection.execute( """ SELECT unit_type, status, COUNT(*) AS count FROM ingest_work_items WHERE run_id = ? GROUP BY unit_type, status ORDER BY unit_type ASC, status ASC """, (run_id,), ).fetchall() counts: dict[str, dict[str, int]] = {} for row in rows: unit_type = str(row["unit_type"] or "unknown") counts.setdefault(unit_type, {status: 0 for status in INGEST_V2_WORK_ITEM_STATUSES}) counts[unit_type][str(row["status"])] = int(row["count"] or 0) return counts def ingest_v2_artifact_counts(connection: sqlite3.Connection, *, run_id: str) -> dict[str, object]: row = connection.execute( """ SELECT COUNT(*) AS count, MIN(updated_at) AS oldest_updated_at FROM ingest_artifact_sweeps WHERE run_id = ? AND state = 'orphan_pending_sweep' """, (run_id,), ).fetchone() orphan_count = int(row["count"] or 0) if row is not None else 0 oldest_age: int | None = None if row is not None and row["oldest_updated_at"]: parsed = parse_utc_timestamp(row["oldest_updated_at"]) if parsed is not None: oldest_age = max(0, int((datetime.now(timezone.utc) - parsed).total_seconds())) return { "orphan_pending_sweep": orphan_count, "oldest_unswept_age_seconds": oldest_age, } def ingest_v2_lease_health(connection: sqlite3.Connection, row: sqlite3.Row, *, run_id: str) -> dict[str, object]: now = datetime.now(timezone.utc) now_text = format_utc_timestamp(now) stale_cutoff = ingest_v2_stale_lease_cutoff(now) active_prepare_workers = int( connection.execute( """ SELECT COUNT(DISTINCT lease_owner) FROM ingest_work_items WHERE run_id = ? AND status = 'leased' AND lease_owner IS NOT NULL AND lease_expires_at IS NOT NULL AND lease_expires_at > ? AND updated_at > ? """, (run_id, now_text, stale_cutoff), ).fetchone()[0] or 0 ) active_lease_row = connection.execute( """ SELECT MIN(updated_at) AS oldest_updated_at FROM ingest_work_items WHERE run_id = ? AND status IN ('leased', 'committing') AND lease_expires_at IS NOT NULL AND lease_expires_at > ? AND updated_at > ? """, (run_id, now_text, stale_cutoff), ).fetchone() stale_lease_row = connection.execute( """ SELECT MIN( CASE WHEN lease_expires_at <= ? THEN lease_expires_at ELSE updated_at END ) AS oldest_expired_at FROM ingest_work_items WHERE run_id = ? AND status IN ('leased', 'committing') AND lease_expires_at IS NOT NULL AND ( lease_expires_at <= ? OR updated_at <= ? ) """, (now_text, run_id, now_text, stale_cutoff), ).fetchone() def age_seconds(raw_value: object) -> int | None: parsed = parse_utc_timestamp(raw_value) if parsed is None: return None return max(0, int((now - parsed).total_seconds())) return { "oldest_active_lease_age_seconds": age_seconds(active_lease_row["oldest_updated_at"] if active_lease_row else None), "oldest_stale_lease_age_seconds": age_seconds(stale_lease_row["oldest_expired_at"] if stale_lease_row else None), "writer_busy": ingest_v2_lease_is_active( row["committer_lease_expires_at"], row["committer_heartbeat_at"], now=now, ), "active_prepare_workers": active_prepare_workers, "prepare_worker_soft_limit": int(row["prepare_worker_soft_limit"] or DEFAULT_WORKER_BACKGROUND_MAX_PARALLEL), } def ingest_v2_next_commands( root: Path, row: sqlite3.Row, *, counts: dict[str, int], leases: dict[str, object], artifacts: dict[str, object], budget_seconds: int, ) -> list[str]: if row["cancel_requested_at"] is not None or str(row["status"]) in INGEST_V2_TERMINAL_STATUSES: return [f"ingest-status {shlex.quote(str(root))} --run-id {shlex.quote(str(row['run_id']))}"] run_id_arg = shlex.quote(str(row["run_id"])) root_arg = shlex.quote(str(root)) budget_arg = str(int(budget_seconds)) runnable_commands: list[str] = [] phase = str(row["phase"] or "") if phase == "planning": runnable_commands.append(f"ingest-plan-step {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}") if ( phase == "preparing" and (counts.get("pending", 0) > 0 or counts.get("leased", 0) > 0) and int(leases.get("active_prepare_workers") or 0) < int(leases.get("prepare_worker_soft_limit") or 0) ): runnable_commands.append(f"ingest-prepare-step {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}") if phase == "committing" and counts.get("prepared", 0) > 0 and not bool(leases.get("writer_busy")): runnable_commands.append(f"ingest-commit-step {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}") if phase == "finalizing" or int(artifacts.get("orphan_pending_sweep") or 0) > 0: runnable_commands.append(f"ingest-finalize-step {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}") commands: list[str] = [] if runnable_commands: commands.append(f"ingest-run-step {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}") commands.extend(runnable_commands) if leases.get("oldest_stale_lease_age_seconds") is not None: commands.append(f"ingest-status {root_arg} --run-id {run_id_arg}") return commands def ingest_v2_status_payload( connection: sqlite3.Connection, root: Path, row: sqlite3.Row, *, budget_seconds: int = DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, ) -> dict[str, object]: run_id = str(row["run_id"]) scope = decode_json_text(row["scope_json"], default={}) or {} counts = ingest_v2_status_counts(connection, run_id=run_id) unit_type_counts = ingest_v2_unit_type_counts(connection, run_id=run_id) leases = ingest_v2_lease_health(connection, row, run_id=run_id) artifacts = ingest_v2_artifact_counts(connection, run_id=run_id) prepared_order_row = connection.execute( """ SELECT MIN(commit_order) AS next_commit_order FROM ingest_work_items WHERE run_id = ? AND status = 'prepared' """, (run_id,), ).fetchone() next_commands = ingest_v2_next_commands( root, row, counts=counts, leases=leases, artifacts=artifacts, budget_seconds=budget_seconds, ) return { "run_id": run_id, "status": row["status"], "phase": row["phase"], "scope": list((scope if isinstance(scope, dict) else {}).get("scan_paths") or []), "scope_details": scope, "budget_recommendation_seconds": int(budget_seconds), "counts": { "work_items": counts, "by_unit_type": unit_type_counts, }, "leases": leases, "entity": { "graph_stale": bool(row["entity_graph_stale"]), "sync_pending": 0, "sync_completed": 0, "entity_writer_busy": False, "policy_changed_during_run": False, }, "artifacts": artifacts, "progress": { "planning_complete": str(row["phase"]) != "planning", "commit_order_next": ( int(prepared_order_row["next_commit_order"]) if prepared_order_row is not None and prepared_order_row["next_commit_order"] is not None else None ), "finalize_phase": str(row["phase"]) if str(row["phase"]) == "finalizing" else None, }, "stalled": leases.get("oldest_stale_lease_age_seconds") is not None, "last_error": row["error"], "next_recommended_commands": next_commands, } def ingest_v2_status_payload_timed( connection: sqlite3.Connection, root: Path, row: sqlite3.Row, *, budget_seconds: int = DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, ) -> tuple[dict[str, object], float]: started = time.perf_counter() payload = ingest_v2_status_payload(connection, root, row, budget_seconds=budget_seconds) return payload, ingest_v2_elapsed_ms(started) def ingest_v2_refresh_conversation_previews_best_effort( connection: sqlite3.Connection, paths: dict[str, Path], conversation_ids: list[int] | None = None, ) -> tuple[int, str | None]: connection.execute("SAVEPOINT ingest_v2_conversation_preview_refresh") try: previews_refreshed = refresh_conversation_previews(connection, paths, conversation_ids) connection.execute("RELEASE SAVEPOINT ingest_v2_conversation_preview_refresh") return previews_refreshed, None except PermissionError as exc: connection.execute("ROLLBACK TO SAVEPOINT ingest_v2_conversation_preview_refresh") connection.execute("RELEASE SAVEPOINT ingest_v2_conversation_preview_refresh") return 0, f"{type(exc).__name__}: {exc}" def ingest_v2_conversation_preview_work_summary( connection: sqlite3.Connection, *, run_id: str, ) -> dict[str, int]: rows = connection.execute( """ SELECT status, artifact_manifest_json FROM ingest_work_items WHERE run_id = ? AND unit_type = 'conversation_preview' """, (run_id,), ).fetchall() summary = { "total": 0, "committed": 0, "failed": 0, "refreshed": 0, } for row in rows: summary["total"] += 1 status = str(row["status"] or "") if status == "committed": summary["committed"] += 1 manifest = decode_json_text(row["artifact_manifest_json"], default={}) or {} if isinstance(manifest, dict): summary["refreshed"] += int(manifest.get("refreshed") or 0) elif status == "failed": summary["failed"] += 1 return summary def ingest_v2_start( root: Path, *, recursive: bool, raw_file_types: str | None, raw_paths: list[str] | None = None, budget_seconds: int | None = None, skip_unchanged_loose_files: bool = True, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) scope = ingest_v2_scope_payload( root, recursive=recursive, raw_file_types=raw_file_types, raw_paths=raw_paths, skip_unchanged_loose_files=skip_unchanged_loose_files, ) with workspace_ingest_session(paths, command_name="ingest-start"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) ensure_workspace_runtime_metadata(root, connection) active_row = active_ingest_v2_run_row(connection) if active_row is not None: raise RetrieverStructuredError( f"Resumable ingest run {active_row['run_id']} is already active.", ingest_v2_conflict_payload( root, active_row, message="A resumable ingest run is active in this workspace.", ), ) run_id = new_ingest_v2_run_id() now = utc_now() connection.execute("BEGIN") try: connection.execute( """ INSERT INTO ingest_runs ( run_id, scope_json, recursive, raw_file_types, pipeline_schema_version, phase, status, prepare_worker_soft_limit, entity_policy_snapshot_json, created_at, started_at, last_heartbeat_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, compact_json_text(scope), 1 if recursive else 0, raw_file_types, INGEST_V2_PIPELINE_SCHEMA_VERSION, "planning", "planning", DEFAULT_WORKER_BACKGROUND_MAX_PARALLEL, "{}", now, now, now, ), ) connection.commit() except Exception: connection.rollback() raise row = require_ingest_v2_run_row(connection, run_id) return { "ok": True, "created": True, **ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } finally: connection.close() def ingest_v2_status(root: Path, *, run_id: str | None = None, budget_seconds: int | None = None) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) if run_id else latest_ingest_v2_run_row(connection) if row is None: return { "ok": True, "status": "none", "phase": None, "run_id": None, "counts": {"work_items": {status: 0 for status in INGEST_V2_WORK_ITEM_STATUSES}, "by_unit_type": {}}, "leases": { "oldest_active_lease_age_seconds": None, "oldest_stale_lease_age_seconds": None, "writer_busy": False, "active_prepare_workers": 0, "prepare_worker_soft_limit": DEFAULT_WORKER_BACKGROUND_MAX_PARALLEL, }, "entity": { "graph_stale": False, "sync_pending": 0, "sync_completed": 0, "entity_writer_busy": False, "policy_changed_during_run": False, }, "artifacts": { "orphan_pending_sweep": 0, "oldest_unswept_age_seconds": None, }, "progress": { "planning_complete": False, "commit_order_next": None, "finalize_phase": None, }, "stalled": False, "last_error": None, "next_recommended_commands": [], } return {"ok": True, **ingest_v2_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() def ingest_v2_cancel(root: Path, *, run_id: str, force: bool = False) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) ensure_layout(paths) with workspace_ingest_session(paths, command_name="ingest-cancel"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) now = utc_now() connection.execute("BEGIN") try: connection.execute( """ UPDATE ingest_runs SET cancel_requested_at = COALESCE(cancel_requested_at, ?), status = CASE WHEN status IN ('completed', 'failed') THEN status ELSE 'canceled' END, phase = CASE WHEN status IN ('completed', 'failed') THEN phase ELSE 'canceled' END, completed_at = COALESCE(completed_at, ?), last_heartbeat_at = ? WHERE run_id = ? """, (now, now, now, run_id), ) cursor = connection.execute( """ UPDATE ingest_work_items SET status = 'cancelled', updated_at = ?, last_error = COALESCE(last_error, 'Ingest run canceled.') WHERE run_id = ? AND status IN ('pending', 'leased', 'prepared') """, (now, run_id), ) cancelled_items = int(cursor.rowcount or 0) ingest_v2_delete_terminal_prepared_items(connection, run_id=run_id) connection.commit() except Exception: connection.rollback() raise ingest_v2_best_effort_cleanup_prepared_payload_spills(paths, run_id=run_id) ingest_v2_best_effort_cleanup_pst_plan_payload_spills(paths, run_id=run_id) updated_row = require_ingest_v2_run_row(connection, run_id) return { "ok": True, "cancel_requested": True, "force": bool(force), "cancelled_work_items": cancelled_items, **ingest_v2_status_payload(connection, root, updated_row), } finally: connection.close() def ingest_v2_plan_step( root: Path, *, run_id: str, budget_seconds: int | None = None, schema_applied: bool = False, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) processed_paths = 0 planned_loose_files = 0 planned_production_roots = 0 planned_production_rows = 0 unsaved_plan_steps = 0 worker_id = ingest_v2_worker_id("plan") cursor_save_ms_values: list[float] = [] work_item_insert_ms_values: list[float] = [] mbox_plan_ms_values: list[float] = [] mbox_plan_payload_bytes = 0 mbox_plan_inline_payload_bytes = 0 mbox_plan_parser_kind_counts: dict[str, int] = {} pst_plan_ms_values: list[float] = [] pst_plan_payload_bytes = 0 pst_plan_spilled_payload_bytes = 0 pst_plan_spilled_messages = 0 pst_plan_attachment_count = 0 pst_plan_attachment_bytes = 0 with workspace_ingest_session(paths, command_name="ingest-plan-step"): connection = connect_db(paths["db_path"]) try: if not schema_applied: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) if str(row["status"]) in INGEST_V2_TERMINAL_STATUSES or row["cancel_requested_at"] is not None: return { "ok": True, "implemented": True, "step": "plan", "processed_paths": 0, "planned_loose_files": 0, "planned_production_roots": 0, "planned_production_rows": 0, "more_planning_remaining": False, "more_work_remaining": False, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } if str(row["phase"]) != "planning": return { "ok": True, "implemented": True, "step": "plan", "processed_paths": 0, "planned_loose_files": 0, "planned_production_roots": 0, "planned_production_rows": 0, "more_planning_remaining": False, "more_work_remaining": str(row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } cursor = ingest_v2_load_or_create_loose_file_plan_cursor(connection, root, row) recursive = bool(row["recursive"]) allowed_types = parse_file_types(row["raw_file_types"]) scan_scope = ingest_v2_scan_scope_from_run(root, row) scope_payload = decode_json_text(row["scope_json"], default={}) or {} skip_unchanged_loose_files = True if isinstance(scope_payload, dict): skip_unchanged_loose_files = bool(scope_payload.get("skip_unchanged_loose_files", True)) while time.perf_counter() < deadline: current_mbox_source = cursor.get("current_mbox_source") if isinstance(current_mbox_source, dict): if unsaved_plan_steps: cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) unsaved_plan_steps = 0 try: updated_mbox_source, next_commit_order, planned_messages, source_complete, plan_metrics = ( ingest_v2_plan_current_mbox_source( connection, root, run_id=run_id, current_mbox_source=current_mbox_source, deadline=deadline, ) ) except Exception as exc: rollback_open_transaction(connection) failures = list(cursor.get("mbox_failures") or []) failures.append( { "source_rel_path": str(current_mbox_source.get("source_rel_path") or ""), "error": f"{type(exc).__name__}: {exc}", } ) cursor["mbox_failures"] = failures cursor["current_mbox_source"] = None cursor["next_commit_order"] = int(current_mbox_source.get("next_commit_order") or cursor.get("next_commit_order") or 1) cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) continue mbox_plan_ms_values.append(float(plan_metrics.get("duration_ms") or 0.0)) mbox_plan_payload_bytes += int(plan_metrics.get("payload_bytes") or 0) mbox_plan_inline_payload_bytes += int(plan_metrics.get("inline_payload_bytes") or 0) for parser_kind, parser_count in dict(plan_metrics.get("parser_kind_counts") or {}).items(): normalized_parser_kind = normalize_whitespace(str(parser_kind or "")) or "unknown" mbox_plan_parser_kind_counts[normalized_parser_kind] = ( int(mbox_plan_parser_kind_counts.get(normalized_parser_kind) or 0) + int(parser_count or 0) ) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=worker_id, event_type="plan_mbox_source_batch", phase="plan", duration_ms=float(plan_metrics.get("duration_ms") or 0.0), details={ "source_rel_path": str(plan_metrics.get("source_rel_path") or ""), "source_plan_kind": str(plan_metrics.get("source_plan_kind") or "mbox"), "planned_messages": int(plan_metrics.get("planned_messages") or 0), "payload_bytes": int(plan_metrics.get("payload_bytes") or 0), "inline_payload_bytes": int(plan_metrics.get("inline_payload_bytes") or 0), "parser_kind_counts": dict(plan_metrics.get("parser_kind_counts") or {}), "message_index_start": plan_metrics.get("message_index_start"), "message_index_end": plan_metrics.get("message_index_end"), "source_complete": bool(plan_metrics.get("source_complete")), }, ) cursor["next_commit_order"] = int(next_commit_order) cursor["planned_mbox_messages"] = int(cursor.get("planned_mbox_messages") or 0) + int(planned_messages) cursor["current_mbox_source"] = updated_mbox_source if source_complete: source_rel_path = str(current_mbox_source.get("source_rel_path") or "") planned_sources = list(cursor.get("planned_mbox_sources") or []) if source_rel_path and source_rel_path not in planned_sources: planned_sources.append(source_rel_path) cursor["planned_mbox_sources"] = planned_sources if source_complete or ingest_v2_deadline_remaining_seconds(deadline) < 1.0: cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) elif isinstance(updated_mbox_source, dict): cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) if not source_complete: break continue current_pst_source = cursor.get("current_pst_source") if isinstance(current_pst_source, dict): if unsaved_plan_steps: cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) unsaved_plan_steps = 0 try: updated_pst_source, next_commit_order, planned_messages, source_complete, plan_metrics = ( ingest_v2_plan_current_pst_source( connection, root, run_id=run_id, current_pst_source=current_pst_source, deadline=deadline, ) ) except Exception as exc: rollback_open_transaction(connection) failures = list(cursor.get("pst_failures") or []) failures.append( { "source_rel_path": str(current_pst_source.get("source_rel_path") or ""), "error": f"{type(exc).__name__}: {exc}", } ) cursor["pst_failures"] = failures cursor["current_pst_source"] = None cursor["next_commit_order"] = int(current_pst_source.get("next_commit_order") or cursor.get("next_commit_order") or 1) cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) continue pst_plan_ms_values.append(float(plan_metrics.get("duration_ms") or 0.0)) pst_plan_payload_bytes += int(plan_metrics.get("payload_bytes") or 0) pst_plan_spilled_payload_bytes += int(plan_metrics.get("spilled_payload_bytes") or 0) pst_plan_spilled_messages += int(plan_metrics.get("spilled_messages") or 0) pst_plan_attachment_count += int(plan_metrics.get("attachment_count") or 0) pst_plan_attachment_bytes += int(plan_metrics.get("attachment_bytes") or 0) ingest_v2_record_worker_event( connection, run_id=run_id, worker_id=worker_id, event_type="plan_pst_source_batch", phase="plan", duration_ms=float(plan_metrics.get("duration_ms") or 0.0), details={ "source_rel_path": str(plan_metrics.get("source_rel_path") or ""), "source_plan_kind": str(plan_metrics.get("source_plan_kind") or "pst"), "planned_messages": int(plan_metrics.get("planned_messages") or 0), "payload_bytes": int(plan_metrics.get("payload_bytes") or 0), "spilled_payload_bytes": int(plan_metrics.get("spilled_payload_bytes") or 0), "spilled_messages": int(plan_metrics.get("spilled_messages") or 0), "attachment_count": int(plan_metrics.get("attachment_count") or 0), "attachment_bytes": int(plan_metrics.get("attachment_bytes") or 0), "message_index_start": plan_metrics.get("message_index_start"), "message_index_end": plan_metrics.get("message_index_end"), "source_complete": bool(plan_metrics.get("source_complete")), }, ) cursor["next_commit_order"] = int(next_commit_order) cursor["planned_pst_messages"] = int(cursor.get("planned_pst_messages") or 0) + int(planned_messages) cursor["current_pst_source"] = updated_pst_source if source_complete: source_rel_path = str(current_pst_source.get("source_rel_path") or "") planned_sources = list(cursor.get("planned_pst_sources") or []) if source_rel_path and source_rel_path not in planned_sources: planned_sources.append(source_rel_path) cursor["planned_pst_sources"] = planned_sources if source_complete or ingest_v2_deadline_remaining_seconds(deadline) < 1.0: cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) elif isinstance(updated_pst_source, dict): cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) if not source_complete: break continue pending_slack_export_roots = list(cursor.get("pending_slack_export_roots") or []) if pending_slack_export_roots: slack_rel_root = str(pending_slack_export_roots.pop(0)) cursor["pending_slack_export_roots"] = pending_slack_export_roots slack_payload = dict( dict(cursor.get("slack_export_roots_by_rel_root") or {}).get(slack_rel_root) or {} ) if slack_payload: try: insert_started = time.perf_counter() slack_plan = ingest_v2_plan_slack_export_root( connection, root, run_id=run_id, payload=slack_payload, next_commit_order=int(cursor.get("next_commit_order") or 1), ) work_item_insert_ms_values.append(ingest_v2_elapsed_ms(insert_started)) except Exception as exc: rollback_open_transaction(connection) failures = list(cursor.get("slack_failures") or []) failures.append( { "slack_rel_root": slack_rel_root, "error": f"{type(exc).__name__}: {exc}", } ) cursor["slack_failures"] = failures slack_plan = None if slack_plan is not None: cursor["next_commit_order"] = int(slack_plan["next_commit_order"]) cursor["planned_slack_conversations"] = ( int(cursor.get("planned_slack_conversations") or 0) + int(slack_plan["planned_conversations"] or 0) ) cursor["planned_slack_day_documents"] = ( int(cursor.get("planned_slack_day_documents") or 0) + int(slack_plan["planned_day_documents"] or 0) ) cursor["planned_slack_day_files_scanned"] = ( int(cursor.get("planned_slack_day_files_scanned") or 0) + int(slack_plan["planned_day_files_scanned"] or 0) ) slack_root_payloads = dict(cursor.get("slack_export_roots_by_rel_root") or {}) slack_root_payloads[slack_rel_root] = { **slack_payload, "seen_rel_paths": list(slack_plan["seen_rel_paths"]), "planned_conversations": int(slack_plan["planned_conversations"] or 0), "planned_day_documents": int(slack_plan["planned_day_documents"] or 0), "planned_day_files_scanned": int(slack_plan["planned_day_files_scanned"] or 0), } cursor["slack_export_roots_by_rel_root"] = slack_root_payloads planned_roots = list(cursor.get("planned_slack_export_roots") or []) if slack_rel_root not in planned_roots: planned_roots.append(slack_rel_root) cursor["planned_slack_export_roots"] = planned_roots cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) continue pending_gmail_mbox_sources = list(cursor.get("pending_gmail_mbox_sources") or []) if pending_gmail_mbox_sources: gmail_source_payload = dict(pending_gmail_mbox_sources.pop(0)) cursor["pending_gmail_mbox_sources"] = pending_gmail_mbox_sources source_rel_path = str(gmail_source_payload.get("source_rel_path") or "") if not source_rel_path: continue scanned_mbox_paths = list(cursor.get("scanned_mbox_source_rel_paths") or []) if source_rel_path not in scanned_mbox_paths: scanned_mbox_paths.append(source_rel_path) cursor["scanned_mbox_source_rel_paths"] = scanned_mbox_paths insert_started = time.perf_counter() source_scan_hash = ingest_v2_gmail_mbox_source_scan_hash( ingest_v2_cursor_path(root, source_rel_path), gmail_source_payload, ) current_mbox_source, next_commit_order, skipped_source = ingest_v2_begin_mbox_source_plan( connection, root, run_id=run_id, rel_path=source_rel_path, commit_order=int(cursor.get("next_commit_order") or 1), source_plan_kind="gmail", source_scan_hash=source_scan_hash, gmail_source_payload=gmail_source_payload, ) work_item_insert_ms_values.append(ingest_v2_elapsed_ms(insert_started)) cursor["next_commit_order"] = int(next_commit_order) cursor["planned_gmail_mbox_sources"] = int(cursor.get("planned_gmail_mbox_sources") or 0) + 1 if skipped_source: cursor["skipped_mbox_sources"] = int(cursor.get("skipped_mbox_sources") or 0) + 1 planned_sources = list(cursor.get("planned_mbox_sources") or []) if source_rel_path not in planned_sources: planned_sources.append(source_rel_path) cursor["planned_mbox_sources"] = planned_sources else: cursor["current_mbox_source"] = current_mbox_source cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) continue pending_production_rel_roots = list(cursor.get("pending_production_rel_roots") or []) if pending_production_rel_roots: production_rel_root = str(pending_production_rel_roots.pop(0)) cursor["pending_production_rel_roots"] = pending_production_rel_roots signature_payload = dict( dict(cursor.get("production_roots_by_rel_root") or {}).get(production_rel_root) or {} ) if signature_payload: try: insert_started = time.perf_counter() production_plan = ingest_v2_plan_production_root( connection, root, run_id=run_id, signature_payload=signature_payload, next_commit_order=int(cursor.get("next_commit_order") or 1), paths=paths, ) work_item_insert_ms_values.append(ingest_v2_elapsed_ms(insert_started)) except Exception as exc: rollback_open_transaction(connection) failures = list(cursor.get("production_failures") or []) failures.append( { "production_rel_root": production_rel_root, "error": f"{type(exc).__name__}: {exc}", } ) cursor["production_failures"] = failures production_plan = None if production_plan is None: cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) continue cursor["next_commit_order"] = int(production_plan["next_commit_order"]) cursor["planned_production_rows"] = ( int(cursor.get("planned_production_rows") or 0) + int(production_plan["planned_rows"] or 0) ) cursor["planned_production_preview_batches"] = ( int(cursor.get("planned_production_preview_batches") or 0) + int(production_plan.get("planned_preview_batches") or 0) ) cursor["skipped_unchanged_production_rows"] = ( int(cursor.get("skipped_unchanged_production_rows") or 0) + int(production_plan.get("skipped_unchanged_rows") or 0) ) cursor["production_docs_missing_linked_text"] = ( int(cursor.get("production_docs_missing_linked_text") or 0) + int(production_plan["docs_missing_linked_text"] or 0) ) cursor["production_docs_missing_linked_images"] = ( int(cursor.get("production_docs_missing_linked_images") or 0) + int(production_plan["docs_missing_linked_images"] or 0) ) cursor["production_docs_missing_linked_natives"] = ( int(cursor.get("production_docs_missing_linked_natives") or 0) + int(production_plan["docs_missing_linked_natives"] or 0) ) production_root_payloads = dict(cursor.get("production_roots_by_rel_root") or {}) production_root_payloads[production_rel_root] = { **signature_payload, "seen_control_numbers": list(production_plan["seen_control_numbers"]), "planned_rows": int(production_plan["planned_rows"] or 0), "planned_preview_batches": int(production_plan.get("planned_preview_batches") or 0), "skipped_unchanged_rows": int(production_plan.get("skipped_unchanged_rows") or 0), "docs_missing_linked_text": int(production_plan["docs_missing_linked_text"] or 0), "docs_missing_linked_images": int(production_plan["docs_missing_linked_images"] or 0), "docs_missing_linked_natives": int(production_plan["docs_missing_linked_natives"] or 0), } cursor["production_roots_by_rel_root"] = production_root_payloads planned_roots = list(cursor.get("planned_production_roots") or []) if production_rel_root not in planned_roots: planned_roots.append(production_rel_root) cursor["planned_production_roots"] = planned_roots planned_production_roots += 1 planned_production_rows += int(production_plan["planned_rows"] or 0) cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) continue pending_paths = list(cursor.get("pending_paths") or []) if not pending_paths: break rel_path = str(pending_paths.pop(0)) cursor["pending_paths"] = pending_paths cursor["scanned_paths"] = int(cursor.get("scanned_paths") or 0) + 1 processed_paths += 1 candidate_path = ingest_v2_cursor_path(root, rel_path) if rel_path and ingest_v2_rel_path_excluded(cursor, rel_path): cursor["skipped_excluded_paths"] = int(cursor.get("skipped_excluded_paths") or 0) + 1 elif not candidate_path.exists(): cursor["skipped_missing_paths"] = int(cursor.get("skipped_missing_paths") or 0) + 1 elif not path_is_at_or_under(candidate_path, root): cursor["skipped_excluded_paths"] = int(cursor.get("skipped_excluded_paths") or 0) + 1 elif candidate_path.is_dir(): child_paths = ingest_v2_planning_child_paths(root, candidate_path, recursive=recursive) cursor["pending_paths"] = ingest_v2_sorted_pending_paths( [*list(cursor.get("pending_paths") or []), *child_paths] ) cursor["listed_directories"] = int(cursor.get("listed_directories") or 0) + 1 elif candidate_path.is_file(): if ingest_v2_path_is_workspace_internal(root, candidate_path): cursor["skipped_excluded_paths"] = int(cursor.get("skipped_excluded_paths") or 0) + 1 elif not ingest_scan_scope_contains_rel_path(scan_scope, rel_path): cursor["skipped_excluded_paths"] = int(cursor.get("skipped_excluded_paths") or 0) + 1 else: file_type = normalize_extension(candidate_path) if not file_type: cursor["skipped_extensionless_files"] = int(cursor.get("skipped_extensionless_files") or 0) + 1 elif allowed_types is not None and file_type not in allowed_types: cursor["skipped_filtered_files"] = int(cursor.get("skipped_filtered_files") or 0) + 1 elif file_type == PST_SOURCE_KIND: scanned_pst_paths = list(cursor.get("scanned_pst_source_rel_paths") or []) if rel_path not in scanned_pst_paths: scanned_pst_paths.append(rel_path) cursor["scanned_pst_source_rel_paths"] = scanned_pst_paths pst_source_payload = dict( dict(cursor.get("pst_source_payloads_by_rel_path") or {}).get(rel_path) or {} ) insert_started = time.perf_counter() current_pst_source, next_commit_order, skipped_source = ingest_v2_begin_pst_source_plan( connection, root, run_id=run_id, rel_path=rel_path, commit_order=int(cursor.get("next_commit_order") or 1), source_payload=pst_source_payload, ) work_item_insert_ms_values.append(ingest_v2_elapsed_ms(insert_started)) cursor["next_commit_order"] = int(next_commit_order) if skipped_source: cursor["skipped_pst_sources"] = int(cursor.get("skipped_pst_sources") or 0) + 1 planned_sources = list(cursor.get("planned_pst_sources") or []) if rel_path not in planned_sources: planned_sources.append(rel_path) cursor["planned_pst_sources"] = planned_sources else: cursor["current_pst_source"] = current_pst_source elif file_type == MBOX_SOURCE_KIND: scanned_mbox_paths = list(cursor.get("scanned_mbox_source_rel_paths") or []) if rel_path not in scanned_mbox_paths: scanned_mbox_paths.append(rel_path) cursor["scanned_mbox_source_rel_paths"] = scanned_mbox_paths insert_started = time.perf_counter() current_mbox_source, next_commit_order, skipped_source = ingest_v2_begin_mbox_source_plan( connection, root, run_id=run_id, rel_path=rel_path, commit_order=int(cursor.get("next_commit_order") or 1), ) work_item_insert_ms_values.append(ingest_v2_elapsed_ms(insert_started)) cursor["next_commit_order"] = int(next_commit_order) if skipped_source: cursor["skipped_mbox_sources"] = int(cursor.get("skipped_mbox_sources") or 0) + 1 planned_sources = list(cursor.get("planned_mbox_sources") or []) if rel_path not in planned_sources: planned_sources.append(rel_path) cursor["planned_mbox_sources"] = planned_sources else: cursor["current_mbox_source"] = current_mbox_source elif file_type in CONTAINER_SOURCE_FILE_TYPES: cursor["skipped_container_files"] = int(cursor.get("skipped_container_files") or 0) + 1 else: file_size, file_mtime_ns = source_file_snapshot(candidate_path) if skip_unchanged_loose_files and ingest_v2_loose_file_matches_existing_snapshot( connection, candidate_path, rel_path=rel_path, file_size=file_size, ): cursor["skipped_unchanged_loose_files"] = ( int(cursor.get("skipped_unchanged_loose_files") or 0) + 1 ) skipped_rel_paths = list(cursor.get("skipped_unchanged_loose_file_rel_paths") or []) if rel_path not in skipped_rel_paths: skipped_rel_paths.append(rel_path) cursor["skipped_unchanged_loose_file_rel_paths"] = skipped_rel_paths else: insert_started = time.perf_counter() inserted = ingest_v2_plan_loose_file_item( connection, run_id=run_id, rel_path=rel_path, file_type=file_type, file_size=file_size, file_mtime_ns=file_mtime_ns, commit_order=int(cursor.get("next_commit_order") or 1), ) work_item_insert_ms_values.append(ingest_v2_elapsed_ms(insert_started)) cursor["next_commit_order"] = int(cursor.get("next_commit_order") or 1) + 1 if inserted: planned_loose_files += 1 cursor["planned_loose_files"] = int(cursor.get("planned_loose_files") or 0) + 1 unsaved_plan_steps += 1 if ( unsaved_plan_steps >= INGEST_V2_PLAN_CURSOR_SAVE_INTERVAL or ingest_v2_deadline_remaining_seconds(deadline) < 1.0 ): cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) unsaved_plan_steps = 0 planning_complete = ( not list(cursor.get("pending_production_rel_roots") or []) and not list(cursor.get("pending_slack_export_roots") or []) and not list(cursor.get("pending_gmail_mbox_sources") or []) and not list(cursor.get("pending_paths") or []) and not isinstance(cursor.get("current_mbox_source"), dict) and not isinstance(cursor.get("current_pst_source"), dict) ) if not planning_complete and unsaved_plan_steps: cursor_save_ms_values.append( ingest_v2_save_planning_cursor_heartbeat( connection, run_id=run_id, cursor=cursor, status="pending", ) ) unsaved_plan_steps = 0 if planning_complete: total_work_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? """, (run_id,), ).fetchone()[0] or 0 ) next_phase = "preparing" if total_work_items else "finalizing" save_started = time.perf_counter() if not connection.in_transaction: connection.execute("BEGIN") try: ingest_v2_save_phase_cursor( connection, run_id=run_id, phase="planning", cursor_key="loose_file_scan", cursor=cursor, status="complete", ) connection.execute( """ UPDATE ingest_runs SET phase = ?, status = ?, last_heartbeat_at = ? WHERE run_id = ? AND cancel_requested_at IS NULL """, (next_phase, next_phase, utc_now(), run_id), ) connection.commit() cursor_save_ms_values.append(ingest_v2_elapsed_ms(save_started)) unsaved_plan_steps = 0 except Exception: connection.rollback() raise updated_row = require_ingest_v2_run_row(connection, run_id) run_payload, status_payload_ms = ingest_v2_status_payload_timed( connection, root, updated_row, budget_seconds=budget, ) return { "ok": True, "implemented": True, "step": "plan", "processed_paths": processed_paths, "planned_loose_files": planned_loose_files, "planned_production_roots": planned_production_roots, "planned_production_rows": planned_production_rows, "cursor": { "pending_paths": len(list(cursor.get("pending_paths") or [])), "pending_production_roots": len(list(cursor.get("pending_production_rel_roots") or [])), "pending_slack_export_roots": len(list(cursor.get("pending_slack_export_roots") or [])), "pending_gmail_mbox_sources": len(list(cursor.get("pending_gmail_mbox_sources") or [])), "scanned_paths": int(cursor.get("scanned_paths") or 0), "planned_loose_files": int(cursor.get("planned_loose_files") or 0), "skipped_unchanged_loose_files": int(cursor.get("skipped_unchanged_loose_files") or 0), "planned_production_roots": len(list(cursor.get("planned_production_roots") or [])), "planned_production_rows": int(cursor.get("planned_production_rows") or 0), "planned_production_preview_batches": int(cursor.get("planned_production_preview_batches") or 0), "skipped_unchanged_production_rows": int(cursor.get("skipped_unchanged_production_rows") or 0), "planned_slack_export_roots": list(cursor.get("planned_slack_export_roots") or []), "planned_slack_conversations": int(cursor.get("planned_slack_conversations") or 0), "planned_slack_day_documents": int(cursor.get("planned_slack_day_documents") or 0), "planned_slack_day_files_scanned": int(cursor.get("planned_slack_day_files_scanned") or 0), "slack_failures": list(cursor.get("slack_failures") or []), "current_mbox_source": ( str(dict(cursor.get("current_mbox_source") or {}).get("source_rel_path") or "") if isinstance(cursor.get("current_mbox_source"), dict) else None ), "planned_mbox_sources": list(cursor.get("planned_mbox_sources") or []), "planned_mbox_messages": int(cursor.get("planned_mbox_messages") or 0), "planned_gmail_mbox_sources": int(cursor.get("planned_gmail_mbox_sources") or 0), "skipped_mbox_sources": int(cursor.get("skipped_mbox_sources") or 0), "scanned_mbox_source_rel_paths": list(cursor.get("scanned_mbox_source_rel_paths") or []), "mbox_failures": list(cursor.get("mbox_failures") or []), "current_pst_source": ( str(dict(cursor.get("current_pst_source") or {}).get("source_rel_path") or "") if isinstance(cursor.get("current_pst_source"), dict) else None ), "planned_pst_sources": list(cursor.get("planned_pst_sources") or []), "planned_pst_messages": int(cursor.get("planned_pst_messages") or 0), "skipped_pst_sources": int(cursor.get("skipped_pst_sources") or 0), "scanned_pst_source_rel_paths": list(cursor.get("scanned_pst_source_rel_paths") or []), "pst_failures": list(cursor.get("pst_failures") or []), "skipped_production_roots": list(cursor.get("skipped_production_roots") or []), "production_failures": list(cursor.get("production_failures") or []), "production_docs_missing_linked_text": int(cursor.get("production_docs_missing_linked_text") or 0), "production_docs_missing_linked_images": int(cursor.get("production_docs_missing_linked_images") or 0), "production_docs_missing_linked_natives": int(cursor.get("production_docs_missing_linked_natives") or 0), "skipped_container_files": int(cursor.get("skipped_container_files") or 0), "skipped_filtered_files": int(cursor.get("skipped_filtered_files") or 0), "skipped_excluded_paths": int(cursor.get("skipped_excluded_paths") or 0), "special_source_counts": cursor.get("special_source_counts") or {}, }, "more_planning_remaining": not planning_complete, "more_work_remaining": str(updated_row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "timings": { "work_item_insert_ms": ingest_v2_timing_summary(work_item_insert_ms_values), "cursor_save_ms": ingest_v2_timing_summary(cursor_save_ms_values), "mbox_source_plan_ms": ingest_v2_timing_summary(mbox_plan_ms_values), "mbox_plan_payload_bytes": int(mbox_plan_payload_bytes), "mbox_plan_inline_payload_bytes": int(mbox_plan_inline_payload_bytes), "mbox_plan_parser_kind_counts": dict(sorted(mbox_plan_parser_kind_counts.items())), "pst_source_plan_ms": ingest_v2_timing_summary(pst_plan_ms_values), "pst_plan_payload_bytes": int(pst_plan_payload_bytes), "pst_plan_spilled_payload_bytes": int(pst_plan_spilled_payload_bytes), "pst_plan_spilled_messages": int(pst_plan_spilled_messages), "pst_plan_attachment_count": int(pst_plan_attachment_count), "pst_plan_attachment_bytes": int(pst_plan_attachment_bytes), "status_payload_ms": round(status_payload_ms, 3), }, "run": run_payload, } finally: connection.close() def ingest_v2_prepare_step( root: Path, *, run_id: str, budget_seconds: int | None = None, schema_applied: bool = False, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) worker_id = ingest_v2_worker_id("prepare") claimed = 0 prepared = 0 deferred_timeout = 0 released = 0 stale_reclaimed = 0 throttled = False claim_limit = ingest_v2_prepare_claim_limit() prepare_workers = 1 prepared_entries: list[dict[str, object]] = [] prepare_ms_values: list[float] = [] prepare_hash_ms_values: list[float] = [] prepare_extract_ms_values: list[float] = [] prepare_chunk_ms_values: list[float] = [] prepared_store_ms_values: list[float] = [] prepared_serialize_ms_values: list[float] = [] prepared_spill_write_ms_values: list[float] = [] prepared_write_ms_values: list[float] = [] prepared_payload_bytes = 0 prepared_spilled_items = 0 prepared_spilled_bytes = 0 connection = connect_db(paths["db_path"]) try: if not schema_applied: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) if str(row["status"]) in INGEST_V2_TERMINAL_STATUSES or row["cancel_requested_at"] is not None: return { "ok": True, "implemented": True, "step": "prepare", "worker_id": worker_id, "claimed": 0, "prepared": 0, "deferred_timeout": 0, "released": 0, "stale_reclaimed": 0, "claim_limit": claim_limit, "prepare_workers": 0, "throttled": False, "more_prepare_remaining": False, "more_work_remaining": False, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } if str(row["phase"]) != "preparing": return { "ok": True, "implemented": True, "step": "prepare", "worker_id": worker_id, "claimed": 0, "prepared": 0, "deferred_timeout": 0, "released": 0, "stale_reclaimed": 0, "claim_limit": claim_limit, "prepare_workers": 0, "throttled": False, "more_prepare_remaining": str(row["phase"]) == "planning", "more_work_remaining": str(row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } stale_reclaimed = ingest_v2_reclaim_stale_prepare_items(connection, run_id=run_id) claimed_rows, throttled = ingest_v2_claim_prepare_items( connection, run_id=run_id, worker_id=worker_id, limit=claim_limit, ) claimed = len(claimed_rows) prepare_workers = max(1, min(ingest_v2_prepare_worker_count_for_rows(claimed_rows), claimed or 1)) prepared_results = ingest_v2_prepare_claimed_work_items_parallel( root, claimed_rows, deadline=deadline, prepare_workers=prepare_workers, ) for prepared_result in prepared_results: work_item_id = int(prepared_result["work_item_id"]) prepared_item = prepared_result.get("prepared_item") source_fingerprint = dict(prepared_result.get("source_fingerprint") or {}) defer_message = prepared_result.get("defer_message") payload_kind = str(prepared_result.get("payload_kind") or "loose_file") if prepared_item is None: if defer_message == "Not enough budget remaining to start prepare.": if ingest_v2_release_prepare_item( connection, run_id=run_id, work_item_id=work_item_id, worker_id=worker_id, reason=defer_message, ): released += 1 else: if ingest_v2_mark_prepare_deferred_timeout( connection, run_id=run_id, work_item_id=work_item_id, worker_id=worker_id, message=defer_message or "Prepare could not complete within the step budget.", ): deferred_timeout += 1 continue prepared_item_dict = dict(prepared_item) prepare_ms_values.append(float(prepared_item_dict.get("prepare_ms") or 0.0)) prepare_hash_ms_values.append(float(prepared_item_dict.get("prepare_hash_ms") or 0.0)) prepare_extract_ms_values.append(float(prepared_item_dict.get("prepare_extract_ms") or 0.0)) prepare_chunk_ms_values.append(float(prepared_item_dict.get("prepare_chunk_ms") or 0.0)) prepared_entries.append( { "work_item_id": work_item_id, "payload_kind": payload_kind, "prepared_item": prepared_item_dict, "source_fingerprint": source_fingerprint, } ) store_result = ingest_v2_store_prepared_items_batch( connection, run_id=run_id, worker_id=worker_id, entries=prepared_entries, ) prepared = int(store_result.get("stored") or 0) prepared_store_ms_values.extend(list(store_result.get("store_total_ms_values") or [])) prepared_serialize_ms_values.extend(list(store_result.get("serialize_ms_values") or [])) prepared_spill_write_ms_values.extend(list(store_result.get("spill_write_ms_values") or [])) prepared_write_ms_values.extend(list(store_result.get("prepared_write_ms_values") or [])) prepared_payload_bytes = int(store_result.get("payload_bytes") or 0) prepared_spilled_items = int(store_result.get("spilled_items") or 0) prepared_spilled_bytes = int(store_result.get("spilled_bytes") or 0) advanced_to_commit = ingest_v2_maybe_advance_after_prepare(connection, run_id=run_id) updated_row = require_ingest_v2_run_row(connection, run_id) remaining_prepare_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND status IN ('pending', 'leased') """, (run_id,), ).fetchone()[0] or 0 ) run_payload, status_payload_ms = ingest_v2_status_payload_timed( connection, root, updated_row, budget_seconds=budget, ) return { "ok": True, "implemented": True, "step": "prepare", "worker_id": worker_id, "claimed": claimed, "prepared": prepared, "deferred_timeout": deferred_timeout, "released": released, "stale_reclaimed": stale_reclaimed, "claim_limit": claim_limit, "prepare_workers": prepare_workers if claimed else 0, "throttled": throttled, "advanced_to_commit": advanced_to_commit, "more_prepare_remaining": remaining_prepare_items > 0, "more_work_remaining": str(updated_row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "timings": { "prepare_ms": ingest_v2_timing_summary(prepare_ms_values), "hash_ms": ingest_v2_timing_summary(prepare_hash_ms_values), "extract_ms": ingest_v2_timing_summary(prepare_extract_ms_values), "chunk_ms": ingest_v2_timing_summary(prepare_chunk_ms_values), "prepared_store_ms": ingest_v2_timing_summary(prepared_store_ms_values), "prepared_serialize_ms": ingest_v2_timing_summary(prepared_serialize_ms_values), "prepared_spill_write_ms": ingest_v2_timing_summary(prepared_spill_write_ms_values), "prepared_write_ms": ingest_v2_timing_summary(prepared_write_ms_values), "prepared_payload_bytes": prepared_payload_bytes, "prepared_spilled_items": prepared_spilled_items, "prepared_spilled_bytes": prepared_spilled_bytes, "status_payload_ms": round(status_payload_ms, 3), }, "run": run_payload, } finally: connection.close() def ingest_v2_commit_step( root: Path, *, run_id: str, budget_seconds: int | None = None, max_items: int | None = None, schema_applied: bool = False, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) item_limit = None if max_items is None else max(1, int(max_items)) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) writer_id = ingest_v2_worker_id("commit") committed = 0 failed = 0 stale_reclaimed = 0 actions: dict[str, int] = {} freshness_fallbacks = 0 writer_busy = False load_commit_state_ms = 0.0 status_payload_ms = 0.0 claim_ms_values: list[float] = [] prepared_decode_ms_values: list[float] = [] item_commit_ms_values: list[float] = [] mbox_message_batch_count = 0 mbox_message_batch_items = 0 mbox_message_batch_fallbacks = 0 pst_message_batch_count = 0 pst_message_batch_items = 0 pst_message_batch_fallbacks = 0 slack_document_batch_count = 0 slack_document_batch_items = 0 slack_document_batch_fallbacks = 0 def commit_timing_payload() -> dict[str, object]: return { "load_commit_state_ms": round(load_commit_state_ms, 3), "claim_ms": ingest_v2_timing_summary(claim_ms_values), "prepared_decode_ms": ingest_v2_timing_summary(prepared_decode_ms_values), "item_commit_ms": ingest_v2_timing_summary(item_commit_ms_values), "status_payload_ms": round(status_payload_ms, 3), } def batching_payload() -> dict[str, int]: return { "mbox_message_batches": int(mbox_message_batch_count), "mbox_message_batch_items": int(mbox_message_batch_items), "mbox_message_batch_fallbacks": int(mbox_message_batch_fallbacks), "pst_message_batches": int(pst_message_batch_count), "pst_message_batch_items": int(pst_message_batch_items), "pst_message_batch_fallbacks": int(pst_message_batch_fallbacks), "slack_document_batches": int(slack_document_batch_count), "slack_document_batch_items": int(slack_document_batch_items), "slack_document_batch_fallbacks": int(slack_document_batch_fallbacks), } connection = connect_db(paths["db_path"]) lease_acquired = False try: if not schema_applied: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) if str(row["status"]) in INGEST_V2_TERMINAL_STATUSES or row["cancel_requested_at"] is not None: run_payload, status_payload_ms = ingest_v2_status_payload_timed( connection, root, row, budget_seconds=budget, ) return { "ok": True, "implemented": True, "step": "commit", "writer_id": writer_id, "writer_busy": False, "committed": 0, "failed": 0, "stale_reclaimed": 0, "actions": {}, "freshness_fallbacks": 0, "advanced_to_finalize": False, "more_commit_remaining": False, "more_work_remaining": False, "batching": batching_payload(), "timings": commit_timing_payload(), "run": run_payload, } if str(row["phase"]) != "committing": run_payload, status_payload_ms = ingest_v2_status_payload_timed( connection, root, row, budget_seconds=budget, ) return { "ok": True, "implemented": True, "step": "commit", "writer_id": writer_id, "writer_busy": False, "committed": 0, "failed": 0, "stale_reclaimed": 0, "actions": {}, "freshness_fallbacks": 0, "advanced_to_finalize": False, "more_commit_remaining": str(row["phase"]) in {"planning", "preparing"}, "more_work_remaining": str(row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "batching": batching_payload(), "timings": commit_timing_payload(), "run": run_payload, } lease_acquired = ingest_v2_acquire_writer_lease(connection, run_id=run_id, writer_id=writer_id) if not lease_acquired: writer_busy = True updated_row = require_ingest_v2_run_row(connection, run_id) run_payload, status_payload_ms = ingest_v2_status_payload_timed( connection, root, updated_row, budget_seconds=budget, ) return { "ok": True, "implemented": True, "step": "commit", "writer_id": writer_id, "writer_busy": True, "committed": 0, "failed": 0, "stale_reclaimed": 0, "actions": {}, "freshness_fallbacks": 0, "advanced_to_finalize": False, "more_commit_remaining": True, "more_work_remaining": str(updated_row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "batching": batching_payload(), "timings": commit_timing_payload(), "run": run_payload, } stale_reclaimed = ingest_v2_reclaim_stale_commit_items(connection, run_id=run_id) run_row = require_ingest_v2_run_row(connection, run_id) load_commit_state_started = time.perf_counter() existing_by_rel, unseen_existing_by_hash = ingest_v2_load_loose_file_commit_state( connection, root=root, run_row=run_row, ) cursor = ingest_v2_load_commit_cursor(connection, run_id=run_id) load_commit_state_ms = ingest_v2_elapsed_ms(load_commit_state_started) filesystem_dataset_id: int | None = None filesystem_dataset_source_id: int | None = None container_contexts: dict[str, dict[str, dict[str, object]]] = { MBOX_SOURCE_KIND: {}, PST_SOURCE_KIND: {}, } def ensure_filesystem_dataset() -> tuple[int, int]: nonlocal filesystem_dataset_id, filesystem_dataset_source_id if filesystem_dataset_id is None or filesystem_dataset_source_id is None: filesystem_dataset_id, filesystem_dataset_source_id = ensure_source_backed_dataset( connection, source_kind=FILESYSTEM_SOURCE_KIND, source_locator=filesystem_dataset_locator(), dataset_name=filesystem_dataset_name(root), ) connection.commit() return filesystem_dataset_id, filesystem_dataset_source_id while ( ingest_v2_deadline_remaining_seconds(deadline) >= INGEST_V2_COMMIT_MIN_START_SECONDS and (item_limit is None or committed + failed < item_limit) ): claim_started = time.perf_counter() claimed_row = ingest_v2_claim_next_commit_item( connection, run_id=run_id, writer_id=writer_id, ) claim_ms_values.append(ingest_v2_elapsed_ms(claim_started)) if claimed_row is None: break work_item_id = int(claimed_row["id"]) item_commit_started = time.perf_counter() timing_item_count = 1 try: prepared_decode_started = time.perf_counter() try: prepared_item = ingest_v2_prepared_item_from_row(claimed_row) finally: prepared_decode_ms_values.append(ingest_v2_elapsed_ms(prepared_decode_started)) payload_kind = str(claimed_row["payload_kind"] or claimed_row["unit_type"] or "") unit_type = str(claimed_row["unit_type"] or "") if payload_kind == "production_row" or unit_type == "production_row": prepared_item = ingest_v2_hydrate_prepared_production_item(prepared_item) dataset_id, dataset_source_id, production_id = ingest_v2_ensure_production_context( connection, root, prepared_item, ) existing_row = connection.execute( """ SELECT * FROM documents WHERE production_id = ? AND control_number = ? ORDER BY id ASC LIMIT 1 """, (production_id, str(prepared_item["control_number"])), ).fetchone() commit_result = commit_prepared_production_row( connection, paths, existing_row, prepared_item, dataset_id=dataset_id, dataset_source_id=dataset_source_id, production_id=production_id, before_transaction_commit=lambda commit_connection, result: ingest_v2_commit_production_work_item_hook( commit_connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, ), ) elif payload_kind == "production_preview_batch" or unit_type == "production_preview_batch": _dataset_id, _dataset_source_id, production_id = ingest_v2_ensure_production_context( connection, root, prepared_item, ) commit_result = ingest_v2_commit_production_preview_batch( connection, paths, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, production_id=production_id, prepared_item=prepared_item, ) elif payload_kind == "slack_conversation" or unit_type == "slack_conversation": prepare_error = normalize_whitespace(str(prepared_item.get("prepare_error") or "")) or None if prepare_error: raise RetrieverError(prepare_error) conversation_identity = prepared_item.get("conversation_identity") identity_root = "" if isinstance(conversation_identity, list) and len(conversation_identity) > 1: identity_root = str(conversation_identity[1]) source_locator = str(prepared_item.get("source_locator") or identity_root) dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=SLACK_EXPORT_SOURCE_KIND, source_locator=source_locator, dataset_name=slack_export_dataset_name(source_locator), ) connection.commit() existing_by_rel_for_slack = existing_rows_by_rel_path( connection, [str(rel_path) for rel_path in list(prepared_item.get("rel_paths") or [])], ) commit_result = commit_prepared_slack_conversation( connection, paths, prepared_item, existing_by_rel_for_slack, dataset_id=int(dataset_id), dataset_source_id=int(dataset_source_id), current_batch=( int(cursor["current_ingestion_batch"]) if cursor.get("current_ingestion_batch") is not None else None ), before_transaction_commit=lambda commit_connection, result: ingest_v2_commit_slack_conversation_work_item_hook( commit_connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, ), ) if str(commit_result.get("status") or "") == "failed": raise RetrieverError(str(commit_result.get("error") or "Slack conversation commit failed.")) elif payload_kind == "slack_document" or unit_type == "slack_document": prepare_error = normalize_whitespace(str(prepared_item.get("prepare_error") or "")) or None if prepare_error: raise RetrieverError(prepare_error) source_locator = str(prepared_item.get("source_locator") or "") dataset_id, dataset_source_id = ensure_source_backed_dataset( connection, source_kind=SLACK_EXPORT_SOURCE_KIND, source_locator=source_locator, dataset_name=slack_export_dataset_name(source_locator), ) connection.commit() batch_claimed_rows = [claimed_row] batch_prepared_items = [prepared_item] additional_rows = ingest_v2_claim_additional_slack_commit_items( connection, run_id=run_id, writer_id=writer_id, source_locator=source_locator, current_payload_bytes=int(claimed_row["payload_bytes"] or 0), ) if additional_rows: additional_prepared_items: list[dict[str, object]] = [] additional_decode_started = time.perf_counter() decode_failed = False for additional_row in additional_rows: try: additional_prepared_items.append(ingest_v2_prepared_item_from_row(additional_row)) except Exception: decode_failed = True break decoded_count = len(additional_prepared_items) if decoded_count: additional_decode_ms = ingest_v2_elapsed_ms(additional_decode_started) prepared_decode_ms_values.extend( [additional_decode_ms / decoded_count] * decoded_count ) if decode_failed: ingest_v2_restore_claimed_commit_items( connection, run_id=run_id, work_item_ids=[int(row["id"]) for row in additional_rows], writer_id=writer_id, ) additional_rows = [] else: batch_claimed_rows.extend(additional_rows) batch_prepared_items.extend(additional_prepared_items) try: batch_results = ingest_v2_commit_slack_document_batch( connection, paths, run_id=run_id, claimed_rows=batch_claimed_rows, prepared_items=batch_prepared_items, writer_id=writer_id, cursor=cursor, dataset_id=int(dataset_id), dataset_source_id=int(dataset_source_id), ) if len(batch_results) > 1: slack_document_batch_count += 1 slack_document_batch_items += len(batch_results) except Exception: rollback_open_transaction(connection) if len(batch_claimed_rows) > 1: slack_document_batch_fallbacks += 1 ingest_v2_restore_claimed_commit_items( connection, run_id=run_id, work_item_ids=[int(row["id"]) for row in batch_claimed_rows[1:]], writer_id=writer_id, ) batch_results = ingest_v2_commit_slack_document_batch( connection, paths, run_id=run_id, claimed_rows=[claimed_row], prepared_items=[prepared_item], writer_id=writer_id, cursor=cursor, dataset_id=int(dataset_id), dataset_source_id=int(dataset_source_id), ) else: raise timing_item_count = max(1, len(batch_results)) for batch_result in batch_results: action = str(batch_result.get("action") or "") if action == "failed": failed += 1 ingest_v2_mark_commit_failed( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, message=str(batch_result.get("error") or "Commit failed."), ) else: committed += 1 actions[action] = int(actions.get(action) or 0) + 1 if bool(batch_result.get("freshness_fallback")): freshness_fallbacks += 1 continue elif payload_kind == "conversation_preview" or unit_type == "conversation_preview": ingest_v2_raise_prepare_error(prepared_item) commit_result = ingest_v2_commit_conversation_preview_work_item( connection, paths, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, prepared_item=prepared_item, ) else: container_work_item = ingest_v2_container_work_item_kind(payload_kind, unit_type) if container_work_item is not None: source_kind, container_role = container_work_item ingest_v2_raise_prepare_error(prepared_item) if container_role == "message": container_batch = ingest_v2_commit_container_message_work_item( connection, paths, run_id=run_id, claimed_row=claimed_row, prepared_item=prepared_item, writer_id=writer_id, cursor=cursor, container_contexts=container_contexts, source_kind=source_kind, prepared_decode_ms_values=prepared_decode_ms_values, ) if source_kind == MBOX_SOURCE_KIND: mbox_message_batch_count += int(container_batch["batch_count"]) mbox_message_batch_items += int(container_batch["batch_items"]) mbox_message_batch_fallbacks += int(container_batch["batch_fallbacks"]) else: pst_message_batch_count += int(container_batch["batch_count"]) pst_message_batch_items += int(container_batch["batch_items"]) pst_message_batch_fallbacks += int(container_batch["batch_fallbacks"]) timing_item_count = int(container_batch["timing_item_count"]) for batch_result in list(container_batch["batch_results"]): action = str(batch_result.get("action") or "") if action == "failed": failed += 1 ingest_v2_mark_commit_failed( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, message=str(batch_result.get("error") or "Commit failed."), ) else: committed += 1 actions[action] = int(actions.get(action) or 0) + 1 continue commit_result = ingest_v2_commit_container_source_finalizer( connection, paths, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, prepared_item=prepared_item, source_kind=source_kind, ) else: commit_result = commit_prepared_loose_file( connection, paths, prepared_item, existing_by_rel, unseen_existing_by_hash, ensure_filesystem_dataset, ( int(cursor["current_ingestion_batch"]) if cursor.get("current_ingestion_batch") is not None else None ), before_transaction_commit=lambda commit_connection, result: ingest_v2_commit_work_item_hook( commit_connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, cursor=cursor, result=result, ), ) action = str(commit_result.get("action") or "") if action == "failed": failed += 1 ingest_v2_mark_commit_failed( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, message=str(commit_result.get("error") or "Commit failed."), ) else: committed += 1 actions[action] = int(actions.get(action) or 0) + 1 if bool(commit_result.get("freshness_fallback")): freshness_fallbacks += 1 except Exception as exc: rollback_open_transaction(connection) failed += 1 ingest_v2_mark_commit_failed( connection, run_id=run_id, work_item_id=work_item_id, writer_id=writer_id, message=f"{type(exc).__name__}: {exc}", ) finally: commit_elapsed_ms = ingest_v2_elapsed_ms(item_commit_started) item_commit_ms_values.extend([commit_elapsed_ms / timing_item_count] * timing_item_count) advanced_after_commit = ingest_v2_maybe_advance_after_commit(connection, run_id=run_id) if lease_acquired: ingest_v2_release_writer_lease(connection, run_id=run_id, writer_id=writer_id) lease_acquired = False updated_row = require_ingest_v2_run_row(connection, run_id) advanced_to_finalize = bool(advanced_after_commit and str(updated_row["phase"]) == "finalizing") advanced_to_prepare = bool(advanced_after_commit and str(updated_row["phase"]) == "preparing") remaining_commit_items = int( connection.execute( """ SELECT COUNT(*) FROM ingest_work_items WHERE run_id = ? AND status IN ('prepared', 'committing') """, (run_id,), ).fetchone()[0] or 0 ) run_payload, status_payload_ms = ingest_v2_status_payload_timed( connection, root, updated_row, budget_seconds=budget, ) return { "ok": True, "implemented": True, "step": "commit", "writer_id": writer_id, "writer_busy": writer_busy, "committed": committed, "failed": failed, "stale_reclaimed": stale_reclaimed, "actions": actions, "freshness_fallbacks": freshness_fallbacks, "advanced_to_finalize": advanced_to_finalize, "advanced_to_prepare": advanced_to_prepare, "more_commit_remaining": remaining_commit_items > 0, "more_work_remaining": str(updated_row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "batching": batching_payload(), "timings": commit_timing_payload(), "run": run_payload, } finally: if lease_acquired: try: ingest_v2_release_writer_lease(connection, run_id=run_id, writer_id=writer_id) except Exception: pass connection.close() def ingest_v2_finalize_step( root: Path, *, run_id: str, budget_seconds: int | None = None, schema_applied: bool = False, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) stages_completed: list[str] = [] connection = connect_db(paths["db_path"]) try: if not schema_applied: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) if str(row["status"]) in INGEST_V2_TERMINAL_STATUSES or row["cancel_requested_at"] is not None: return { "ok": True, "implemented": True, "step": "finalize", "stages_completed": [], "finalization_complete": str(row["status"]) == "completed", "more_finalize_remaining": False, "more_work_remaining": False, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } if str(row["phase"]) != "finalizing": return { "ok": True, "implemented": True, "step": "finalize", "stages_completed": [], "finalization_complete": False, "more_finalize_remaining": str(row["phase"]) in {"planning", "preparing", "committing"}, "more_work_remaining": str(row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } cursor = ingest_v2_load_finalize_cursor(connection, run_id=run_id) stage = str(cursor.get("stage") or "missing") if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_COMMIT_MIN_START_SECONDS: return { "ok": True, "implemented": True, "step": "finalize", "stages_completed": [], "finalization_complete": False, "more_finalize_remaining": True, "more_work_remaining": True, "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } if stage == "production": pending_production_rel_roots = list(cursor.get("pending_production_rel_roots") or []) while ( pending_production_rel_roots and ingest_v2_deadline_remaining_seconds(deadline) >= INGEST_V2_COMMIT_MIN_START_SECONDS ): production_rel_root = str(pending_production_rel_roots.pop(0)) cursor["pending_production_rel_roots"] = pending_production_rel_roots production_payload = dict( dict(cursor.get("production_roots_by_rel_root") or {}).get(production_rel_root) or {} ) if production_payload: production_stats = ingest_v2_finalize_production_root( connection, paths, root, production_payload=production_payload, ) aggregate_stats = cursor.setdefault("production_stats", {}) if isinstance(aggregate_stats, dict): for key, value in production_stats.items(): aggregate_stats[key] = int(aggregate_stats.get(key) or 0) + int(value) finalized_roots = list(cursor.get("production_finalized_roots") or []) finalized_roots.append(production_rel_root) cursor["production_finalized_roots"] = finalized_roots connection.execute("BEGIN") try: ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="pending") connection.commit() except Exception: connection.rollback() raise if pending_production_rel_roots: updated_row = require_ingest_v2_run_row(connection, run_id) return { "ok": True, "implemented": True, "step": "finalize", "stages_completed": stages_completed, "cursor": cursor, "finalization_complete": False, "more_finalize_remaining": True, "more_work_remaining": True, "run": ingest_v2_status_payload(connection, root, updated_row, budget_seconds=budget), } cursor["stage"] = "missing" connection.execute("BEGIN") try: ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="pending") connection.commit() except Exception: connection.rollback() raise if list(cursor.get("production_finalized_roots") or []): stages_completed.append("production") stage = "missing" if stage == "missing": scanned_rel_paths = ingest_v2_run_loose_rel_paths(connection, run_id=run_id) scan_scope = ingest_v2_scan_scope_from_run(root, row) filesystem_missing = mark_missing_documents(connection, scanned_rel_paths, scan_scope=scan_scope) scanned_pst_source_rel_paths = ingest_v2_run_pst_source_rel_paths(connection, run_id=run_id) pst_sources_missing, pst_documents_missing = mark_missing_pst_documents( connection, scanned_pst_source_rel_paths, scan_scope=scan_scope, ) scanned_mbox_source_rel_paths = ingest_v2_run_mbox_source_rel_paths(connection, run_id=run_id) mbox_sources_missing, mbox_documents_missing = mark_missing_mbox_documents( connection, scanned_mbox_source_rel_paths, scan_scope=scan_scope, ) slack_documents_missing = ingest_v2_mark_missing_slack_documents(connection, run_id=run_id) cursor["filesystem_missing"] = int(filesystem_missing) cursor["pst_sources_missing"] = int(pst_sources_missing) cursor["pst_documents_missing"] = int(pst_documents_missing) cursor["mbox_sources_missing"] = int(mbox_sources_missing) cursor["mbox_documents_missing"] = int(mbox_documents_missing) cursor["slack_documents_missing"] = int(slack_documents_missing) cursor["stage"] = "conversations" connection.execute("BEGIN") try: ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="pending") connection.commit() except Exception: connection.rollback() raise stages_completed.append("missing") stage = "conversations" if stage == "conversations" and ingest_v2_deadline_remaining_seconds(deadline) >= INGEST_V2_COMMIT_MIN_START_SECONDS: preview_stage_processed = False connection.execute("BEGIN") try: conversation_assignment = assign_supported_conversations(connection) mbox_documents_redeleted = ingest_v2_redelete_retired_mbox_documents( connection, paths, run_id=run_id, ) pst_documents_redeleted = ingest_v2_redelete_retired_pst_documents( connection, paths, run_id=run_id, ) cursor["conversation_assignment"] = { key: int(value) for key, value in dict(conversation_assignment).items() } cursor["mbox_documents_redeleted"] = int(mbox_documents_redeleted) cursor["pst_documents_redeleted"] = int(pst_documents_redeleted) active_conversation_ids = list_active_conversation_ids(connection) forced_conversation_ids = ingest_v2_run_affected_conversation_ids(connection, run_id=run_id) target_conversation_ids = ingest_v2_conversation_previews_need_refresh( connection, paths, active_conversation_ids, force_conversation_ids=forced_conversation_ids, ) preview_stage_processed = bool( active_conversation_ids or forced_conversation_ids or target_conversation_ids ) cursor["conversation_preview_active_count"] = len(active_conversation_ids) cursor["conversation_preview_target_count"] = len(target_conversation_ids) cursor["conversation_preview_skipped_fresh"] = max( 0, len(active_conversation_ids) - len(target_conversation_ids), ) cursor["conversation_preview_forced_count"] = len( [ conversation_id for conversation_id in target_conversation_ids if conversation_id in forced_conversation_ids ] ) cursor["conversation_preview_work_items_planned"] = 0 previews_refreshed, preview_error = ingest_v2_refresh_conversation_previews_best_effort( connection, paths, target_conversation_ids, ) cursor["conversation_previews_refreshed"] = int(previews_refreshed) cursor["conversation_preview_failures"] = 1 if preview_error else 0 cursor["conversation_preview_last_error"] = preview_error cursor["empty_conversation_preview_dirs_pruned"] = prune_empty_conversation_preview_dirs(paths) cursor["stage"] = "prune" next_phase = "finalizing" ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="pending") connection.execute( """ UPDATE ingest_runs SET phase = ?, status = ?, last_heartbeat_at = ? WHERE run_id = ? AND phase = 'finalizing' AND cancel_requested_at IS NULL """, (next_phase, next_phase, utc_now(), run_id), ) connection.commit() except Exception: connection.rollback() raise stages_completed.append("conversations") if preview_stage_processed: stages_completed.append("conversation_previews") stage = str(cursor.get("stage") or "prune") if ( stage == "conversation_previews" and str(require_ingest_v2_run_row(connection, run_id)["phase"]) == "finalizing" ): preview_summary = ingest_v2_conversation_preview_work_summary(connection, run_id=run_id) connection.execute("BEGIN") try: cursor["conversation_previews_refreshed"] = int(preview_summary["refreshed"]) cursor["conversation_preview_failures"] = int(preview_summary["failed"]) cursor["mbox_documents_redeleted"] = int(cursor.get("mbox_documents_redeleted") or 0) + int( ingest_v2_redelete_retired_mbox_documents( connection, paths, run_id=run_id, ) ) cursor["pst_documents_redeleted"] = int(cursor.get("pst_documents_redeleted") or 0) + int( ingest_v2_redelete_retired_pst_documents( connection, paths, run_id=run_id, ) ) cursor["empty_conversation_preview_dirs_pruned"] = prune_empty_conversation_preview_dirs(paths) cursor["stage"] = "prune" ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="pending") connection.commit() except Exception: connection.rollback() raise stages_completed.append("conversation_previews") stage = "prune" if stage == "prune" and ingest_v2_deadline_remaining_seconds(deadline) >= INGEST_V2_COMMIT_MIN_START_SECONDS: connection.execute("BEGIN") try: pruned = prune_unused_filesystem_dataset(connection) cursor["pruned_unused_filesystem_dataset"] = bool(pruned) cursor["stage"] = "complete" ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="pending") connection.commit() except Exception: connection.rollback() raise stages_completed.append("prune") stage = "complete" finalization_complete = stage == "complete" if finalization_complete: now = utc_now() connection.execute("BEGIN") try: ingest_v2_save_finalize_cursor(connection, run_id=run_id, cursor=cursor, status="complete") connection.execute( """ UPDATE ingest_runs SET phase = 'completed', status = 'completed', completed_at = COALESCE(completed_at, ?), last_heartbeat_at = ?, error = NULL WHERE run_id = ? AND phase = 'finalizing' AND cancel_requested_at IS NULL """, (now, now, run_id), ) connection.commit() except Exception: connection.rollback() raise ingest_v2_delete_terminal_prepared_items(connection, run_id=run_id) connection.commit() ingest_v2_best_effort_cleanup_prepared_payload_spills(paths, run_id=run_id) ingest_v2_best_effort_cleanup_pst_plan_payload_spills(paths, run_id=run_id) ingest_v2_best_effort_sqlite_cleanup(connection) stages_completed.append("complete") updated_row = require_ingest_v2_run_row(connection, run_id) return { "ok": True, "implemented": True, "step": "finalize", "stages_completed": stages_completed, "cursor": cursor, "finalization_complete": str(updated_row["status"]) == "completed", "more_finalize_remaining": str(updated_row["phase"]) == "finalizing", "more_work_remaining": str(updated_row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "run": ingest_v2_status_payload(connection, root, updated_row, budget_seconds=budget), } finally: connection.close() def ingest_v2_select_runnable_step(status_payload: dict[str, object]) -> str | None: for command in list(status_payload.get("next_recommended_commands") or []): try: command_name = shlex.split(str(command))[0] except (IndexError, ValueError): continue if command_name == "ingest-run-step": continue if command_name == "ingest-plan-step": return "plan" if command_name == "ingest-prepare-step": return "prepare" if command_name == "ingest-commit-step": return "commit" if command_name == "ingest-finalize-step": return "finalize" return None def ingest_v2_runner_step_budget(remaining_seconds: float) -> int: return max( 1, min( MAX_RESUMABLE_STEP_BUDGET_SECONDS, int(max(1.0, remaining_seconds)), ), ) def ingest_v2_run_step( root: Path, *, run_id: str | None = None, budget_seconds: int | None = None, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) schema_started = time.perf_counter() schema_ms = 0.0 initial_status_payload_ms = 0.0 stale_reclaimed = {"prepare": 0, "commit": 0} try: apply_schema(connection, root) schema_ms = ingest_v2_elapsed_ms(schema_started) row = require_ingest_v2_run_row(connection, run_id) if run_id else latest_ingest_v2_run_row(connection) if row is None: status_payload = ingest_v2_status(root, run_id=None, budget_seconds=budget) run_payload = dict(status_payload) run_payload.pop("ok", None) return { "ok": True, "implemented": True, "step": "run", "executed": False, "selected_step": None, "executed_steps": [], "stale_reclaimed": stale_reclaimed, "reason": "no_ingest_run", "step_result": None, "step_results": [], "more_work_remaining": False, "run": run_payload, "timings": { "schema_ms": round(schema_ms, 3), "initial_status_payload_ms": 0.0, "inner_step_ms": ingest_v2_timing_summary([]), }, "remaining_budget_seconds": round(ingest_v2_deadline_remaining_seconds(deadline), 3), } resolved_run_id = str(row["run_id"]) if str(row["status"]) in INGEST_V2_ACTIVE_STATUSES and row["cancel_requested_at"] is None: stale_reclaimed = ingest_v2_reclaim_stale_work_items(connection, run_id=resolved_run_id) if any(stale_reclaimed.values()): row = require_ingest_v2_run_row(connection, resolved_run_id) status_payload, initial_status_payload_ms = ingest_v2_status_payload_timed( connection, root, row, budget_seconds=budget, ) finally: connection.close() executed_steps: list[str] = [] step_results: list[dict[str, object]] = [] inner_step_ms_values: list[float] = [] stop_reason: str | None = None run_payload = status_payload while len(executed_steps) < INGEST_V2_RUN_STEP_MAX_INNER_STEPS: remaining_seconds = ingest_v2_deadline_remaining_seconds(deadline) if remaining_seconds < INGEST_V2_RUN_STEP_MIN_REMAINING_SECONDS: stop_reason = "budget_exhausted" break selected_step = ingest_v2_select_runnable_step(run_payload) if selected_step is None: stop_reason = ( "run_terminal" if str(run_payload.get("status")) in INGEST_V2_TERMINAL_STATUSES else "no_runnable_step" ) break inner_budget = ingest_v2_runner_step_budget(remaining_seconds) step_started = time.perf_counter() if selected_step == "plan": step_result = ingest_v2_plan_step( root, run_id=resolved_run_id, budget_seconds=inner_budget, schema_applied=True, ) elif selected_step == "prepare": step_result = ingest_v2_prepare_step( root, run_id=resolved_run_id, budget_seconds=inner_budget, schema_applied=True, ) elif selected_step == "commit": step_result = ingest_v2_commit_step( root, run_id=resolved_run_id, budget_seconds=inner_budget, schema_applied=True, ) elif selected_step == "finalize": step_result = ingest_v2_finalize_step( root, run_id=resolved_run_id, budget_seconds=inner_budget, schema_applied=True, ) else: raise RetrieverError(f"Unsupported resumable ingest step selection: {selected_step}") inner_step_ms_values.append(ingest_v2_elapsed_ms(step_started)) executed_steps.append(selected_step) step_results.append(step_result) if isinstance(step_result.get("run"), dict): run_payload = dict(step_result["run"]) if selected_step == "prepare" and bool(step_result.get("throttled")): stop_reason = "prepare_throttled" break if selected_step == "commit" and bool(step_result.get("writer_busy")): stop_reason = "writer_busy" break if not bool(step_result.get("more_work_remaining")): stop_reason = "run_terminal" break else: stop_reason = "max_inner_steps" return { "ok": True, "implemented": True, "step": "run", "executed": bool(executed_steps), "selected_step": executed_steps[0] if executed_steps else None, "executed_steps": executed_steps, "stale_reclaimed": stale_reclaimed, "reason": stop_reason, "step_result": step_results[-1] if step_results else None, "step_results": step_results, "more_work_remaining": str(run_payload.get("status")) not in INGEST_V2_TERMINAL_STATUSES, "run": run_payload, "timings": { "schema_ms": round(schema_ms, 3), "initial_status_payload_ms": round(initial_status_payload_ms, 3), "inner_step_ms": ingest_v2_timing_summary(inner_step_ms_values), }, "remaining_budget_seconds": round(ingest_v2_deadline_remaining_seconds(deadline), 3), } def ingest_v2_scopes_match(active_scope: object, requested_scope: dict[str, object]) -> bool: if not isinstance(active_scope, dict): return False return compact_json_text(active_scope) == compact_json_text(requested_scope) def ingest_v2_facade_command( root: Path, *, recursive: bool, raw_file_types: str | None, raw_paths: list[str] | None, budget_seconds: int, ) -> str: parts = ["ingest", shlex.quote(str(root))] if recursive: parts.append("--recursive") for raw_path in raw_paths or []: parts.extend(["--path", shlex.quote(str(raw_path))]) if raw_file_types: parts.extend(["--file-types", shlex.quote(str(raw_file_types))]) parts.extend(["--budget-seconds", str(int(budget_seconds))]) return " ".join(parts) def ingest_v2_facade_payload( *, root: Path, recursive: bool, raw_file_types: str | None, raw_paths: list[str] | None, budget_seconds: int, created: bool, resumed: bool, run_payload: dict[str, object], step_payloads: list[dict[str, object]], reason: str, mode: str, ) -> dict[str, object]: run_id = str(run_payload.get("run_id") or "") more_work_remaining = str(run_payload.get("status")) not in INGEST_V2_TERMINAL_STATUSES executed_steps: list[str] = [] for step_payload in step_payloads: executed_steps.extend(str(step) for step in list(step_payload.get("executed_steps") or [])) next_commands: list[str] = [] if more_work_remaining: next_commands.append( ingest_v2_facade_command( root, recursive=recursive, raw_file_types=raw_file_types, raw_paths=raw_paths, budget_seconds=budget_seconds, ) ) next_commands.extend(str(command) for command in list(run_payload.get("next_recommended_commands") or [])) return { "ok": True, "pipeline": INGEST_PIPELINE_V2, "mode": mode, "created": created, "resumed": resumed, "run_id": run_id or None, "status": run_payload.get("status"), "phase": run_payload.get("phase"), "reason": reason, "executed": bool(executed_steps), "executed_steps": executed_steps, "step_calls": len(step_payloads), "step_results": step_payloads, "more_work_remaining": more_work_remaining, "run": run_payload, "counts": run_payload.get("counts"), "next_recommended_commands": next_commands, } def ingest_v2_facade_remaining_budget(deadline: float) -> int | None: remaining_seconds = ingest_v2_deadline_remaining_seconds(deadline) if remaining_seconds < INGEST_V2_RUN_STEP_MIN_REMAINING_SECONDS: return None return max(1, min(MAX_RESUMABLE_STEP_BUDGET_SECONDS, int(remaining_seconds))) def ingest_v2_facade( root: Path, *, recursive: bool, raw_file_types: str | None, raw_paths: list[str] | None = None, budget_seconds: int | None = None, run_to_completion: bool = False, skip_unchanged_loose_files: bool = True, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) requested_scope = ingest_v2_scope_payload( root, recursive=recursive, raw_file_types=raw_file_types, raw_paths=raw_paths, skip_unchanged_loose_files=skip_unchanged_loose_files, ) created = False resumed = False connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) ensure_workspace_runtime_metadata(root, connection) active_row = active_ingest_v2_run_row(connection) if active_row is not None: active_scope = decode_json_text(active_row["scope_json"], default={}) or {} if not ingest_v2_scopes_match(active_scope, requested_scope): payload = ingest_v2_conflict_payload( root, active_row, message="A resumable ingest run is active with a different scope.", ) payload["active_scope"] = active_scope payload["requested_scope"] = requested_scope raise RetrieverStructuredError( f"Resumable ingest run {active_row['run_id']} is already active with a different scope.", payload, ) run_id = str(active_row["run_id"]) resumed = True run_payload = ingest_v2_status_payload(connection, root, active_row, budget_seconds=budget) else: run_id = "" run_payload = {} finally: connection.close() if not resumed: start_payload = ingest_v2_start( root, recursive=recursive, raw_file_types=raw_file_types, raw_paths=raw_paths, budget_seconds=budget, skip_unchanged_loose_files=skip_unchanged_loose_files, ) run_id = str(start_payload["run_id"]) created = True run_payload = { key: value for key, value in start_payload.items() if key not in {"ok", "created"} } if progress_callback is not None and run_to_completion: progress_callback((start_payload if not resumed else run_payload), True) step_payloads: list[dict[str, object]] = [] reason = "run_terminal" if str(run_payload.get("status")) in INGEST_V2_TERMINAL_STATUSES else "budget_exhausted" while str(run_payload.get("status")) not in INGEST_V2_TERMINAL_STATUSES: step_budget = budget if run_to_completion else ingest_v2_facade_remaining_budget(deadline) if step_budget is None: reason = "budget_exhausted" break step_payload = ingest_v2_run_step(root, run_id=run_id, budget_seconds=step_budget) step_payloads.append(step_payload) run_payload = dict(step_payload["run"]) reason = str(step_payload.get("reason") or reason) if progress_callback is not None and run_to_completion: progress_callback(step_payload, False) if not run_to_completion: break payload = ingest_v2_facade_payload( root=root, recursive=recursive, raw_file_types=raw_file_types, raw_paths=raw_paths, budget_seconds=budget, created=created, resumed=resumed, run_payload=run_payload, step_payloads=step_payloads, reason=reason, mode="run_to_completion" if run_to_completion else "bounded", ) if progress_callback is not None and run_to_completion: progress_callback(payload, True) return payload def ingest_v2_step_not_implemented( root: Path, *, run_id: str, phase: str, budget_seconds: int | None = None, ) -> dict[str, object]: set_active_workspace_root(root) budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = require_ingest_v2_run_row(connection, run_id) return { "ok": True, "implemented": False, "step": phase, "processed": 0, "more_work_remaining": str(row["status"]) not in INGEST_V2_TERMINAL_STATUSES, "message": "This V2 step command is reserved by the foundation slice and will be implemented in the next ingest slice.", "run": ingest_v2_status_payload(connection, root, row, budget_seconds=budget), } finally: connection.close() def source_file_snapshot(path: Path) -> tuple[int | None, int | None]: try: stat = path.stat() except FileNotFoundError: return None, None return stat.st_size, stat.st_mtime_ns def rollback_open_transaction(connection: sqlite3.Connection) -> None: if not connection.in_transaction: return try: connection.rollback() except sqlite3.Error: return def refresh_ingest_item_filesystem_facts(item: dict[str, object]) -> dict[str, object]: refreshed_item = dict(item) path = Path(refreshed_item["path"]) file_size, file_mtime_ns = source_file_snapshot(path) refreshed_item["source_file_size"] = file_size refreshed_item["source_file_mtime_ns"] = file_mtime_ns refreshed_item["file_hash"] = ( None if str(refreshed_item["file_type"]) in CONTAINER_SOURCE_FILE_TYPES else sha256_file(path) ) return refreshed_item def plan_ingest_work( root: Path, recursive: bool, allowed_types: set[str] | None, connection: sqlite3.Connection, scan_scope: dict[str, object], ) -> dict[str, object]: scan_hash_ms = 0.0 production_signatures = find_production_root_signatures(root, recursive, connection, scan_scope=scan_scope) production_root_paths = [Path(signature["root"]).resolve() for signature in production_signatures] slack_export_descriptors = find_scoped_source_roots( find_slack_export_roots, root, recursive, scan_scope, allowed_types, ) slack_export_root_paths = [Path(descriptor["root"]).resolve() for descriptor in slack_export_descriptors] gmail_export_descriptors = find_scoped_source_roots( find_gmail_export_roots, root, recursive, scan_scope, allowed_types, ) pst_export_descriptors = ( find_scoped_source_roots(find_pst_export_roots, root, recursive, scan_scope) if allowed_types is None or PST_SOURCE_KIND in allowed_types else [] ) gmail_owned_paths = { Path(path).resolve() for descriptor in gmail_export_descriptors for path in list(descriptor.get("owned_paths") or []) } pst_export_owned_paths = { Path(path).resolve() for descriptor in pst_export_descriptors for path in list(descriptor.get("owned_paths") or []) } gmail_owned_rel_paths = { relative_document_path(root, owned_path) for owned_path in gmail_owned_paths if root.resolve() == owned_path.resolve() or root.resolve() in owned_path.resolve().parents } pst_export_owned_rel_paths = { relative_document_path(root, owned_path) for owned_path in pst_export_owned_paths if root.resolve() == owned_path.resolve() or root.resolve() in owned_path.resolve().parents } pst_export_descriptors_by_pst_path = { Path(path).resolve().as_posix(): descriptor for descriptor in pst_export_descriptors for path in list(descriptor.get("pst_paths") or []) } scanned_files = [ path for path in collect_files(root, recursive, allowed_types, scan_scope=scan_scope) if not any(production_root == path.resolve() or production_root in path.resolve().parents for production_root in production_root_paths) and not any(slack_root == path.resolve() or slack_root in path.resolve().parents for slack_root in slack_export_root_paths) and path.resolve() not in gmail_owned_paths and path.resolve() not in pst_export_owned_paths ] scanned_rel_paths: set[str] = set() scanned_pst_source_rel_paths: set[str] = set() scanned_mbox_source_rel_paths: set[str] = set() scanned_items: list[dict[str, object]] = [] loose_file_items: list[dict[str, object]] = [] for path in scanned_files: item_started = time.perf_counter() rel_path = relative_document_path(root, path) file_type = normalize_extension(path) scanned_rel_paths.add(rel_path) if file_type == PST_SOURCE_KIND: scanned_pst_source_rel_paths.add(rel_path) if file_type == MBOX_SOURCE_KIND: scanned_mbox_source_rel_paths.add(rel_path) item = refresh_ingest_item_filesystem_facts( { "path": path, "rel_path": rel_path, "file_type": file_type, } ) scanned_items.append(item) if file_type not in CONTAINER_SOURCE_FILE_TYPES: loose_file_items.append(item) scan_hash_ms += (time.perf_counter() - item_started) * 1000.0 return { "production_signatures": production_signatures, "slack_export_descriptors": slack_export_descriptors, "gmail_export_descriptors": gmail_export_descriptors, "pst_export_descriptors": pst_export_descriptors, "gmail_owned_rel_paths": gmail_owned_rel_paths, "pst_export_owned_rel_paths": pst_export_owned_rel_paths, "pst_export_descriptors_by_pst_path": pst_export_descriptors_by_pst_path, "scanned_files": scanned_files, "scanned_items": scanned_items, "loose_file_items": loose_file_items, "scanned_rel_paths": scanned_rel_paths, "scanned_pst_source_rel_paths": scanned_pst_source_rel_paths, "scanned_mbox_source_rel_paths": scanned_mbox_source_rel_paths, "scan_hash_ms": scan_hash_ms, } def ingest_serial_special_sources( connection: sqlite3.Connection, paths: dict[str, Path], root: Path, ingest_tmp_dir: Path | None, allowed_types: set[str] | None, production_signatures: list[dict[str, object]], slack_export_descriptors: list[dict[str, object]], gmail_export_descriptors: list[dict[str, object]], scanned_rel_paths: set[str], scanned_mbox_source_rel_paths: set[str], stats: dict[str, int], failures: list[dict[str, str]], ) -> dict[str, object]: current_ingestion_batch: int | None = None slack_day_documents_missing = 0 ingested_production_roots: list[str] = [] skipped_production_roots: list[str] = [] warnings: list[str] = [] gmail_ms = 0.0 for descriptor in gmail_export_descriptors: export_root = Path(descriptor["root"]) source_started = time.perf_counter() try: gmail_result = ingest_gmail_export_root( connection, paths, root, descriptor, allowed_file_types=allowed_types, staging_root=ingest_tmp_dir, ) stats["new"] += int(gmail_result["new"]) stats["updated"] += int(gmail_result["updated"]) stats["failed"] += int(gmail_result["failed"]) stats["gmail_documents_scanned"] += int(gmail_result["scanned_files"]) stats["mbox_sources_skipped"] += int(gmail_result["mbox_sources_skipped"]) stats["mbox_messages_created"] += int(gmail_result["mbox_messages_created"]) stats["mbox_messages_updated"] += int(gmail_result["mbox_messages_updated"]) stats["mbox_messages_deleted"] += int(gmail_result["mbox_messages_deleted"]) stats["gmail_linked_documents_created"] += int(gmail_result["gmail_linked_documents_created"]) stats["gmail_linked_documents_updated"] += int(gmail_result["gmail_linked_documents_updated"]) scanned_rel_paths.update(str(rel_path) for rel_path in list(gmail_result.get("scanned_filesystem_rel_paths", []))) scanned_mbox_source_rel_paths.update( str(rel_path) for rel_path in list(gmail_result.get("scanned_mbox_source_rel_paths", [])) ) failures.extend(list(gmail_result.get("failures", []))) except Exception as exc: rollback_open_transaction(connection) stats["failed"] += 1 failures.append( { "rel_path": relative_document_path(root, export_root), "error": f"{type(exc).__name__}: {exc}", } ) finally: gmail_ms += (time.perf_counter() - source_started) * 1000.0 slack_ms = 0.0 for descriptor in slack_export_descriptors: export_root = Path(descriptor["root"]) source_started = time.perf_counter() try: slack_result = ingest_slack_export_root( connection, paths, export_root, ingestion_batch_number=current_ingestion_batch, staging_root=ingest_tmp_dir, ) current_ingestion_batch = ( int(slack_result["ingestion_batch_number"]) if slack_result.get("ingestion_batch_number") is not None else current_ingestion_batch ) stats["new"] += int(slack_result["new"]) stats["updated"] += int(slack_result["updated"]) stats["failed"] += int(slack_result["failed"]) stats["slack_day_documents_scanned"] += int(slack_result["scanned_day_files"]) stats["slack_documents_created"] += int(slack_result["new"]) stats["slack_documents_updated"] += int(slack_result["updated"]) stats["slack_conversations"] += int(slack_result["conversations"]) slack_day_documents_missing += int(slack_result["missing"]) failures.extend(list(slack_result.get("failures", []))) except Exception as exc: rollback_open_transaction(connection) stats["failed"] += 1 failures.append( { "rel_path": relative_document_path(root, export_root), "error": f"{type(exc).__name__}: {exc}", } ) finally: slack_ms += (time.perf_counter() - source_started) * 1000.0 production_ms = 0.0 if allowed_types is None: for signature in production_signatures: production_rel_root = str(signature["rel_root"]) resolved_production_root = Path(signature["root"]).resolve() source_started = time.perf_counter() try: production_result = ingest_resolved_production_root( connection, paths, root, resolved_production_root, staging_root=ingest_tmp_dir, ) ingested_production_roots.append(production_rel_root) stats["new"] += int(production_result["created"]) stats["updated"] += int(production_result["updated"]) stats["production_documents_created"] += int(production_result["created"]) stats["production_documents_updated"] += int(production_result["updated"]) stats["production_documents_unchanged"] += int(production_result["unchanged"]) stats["production_documents_retired"] += int(production_result["retired"]) stats["production_families_reconstructed"] += int(production_result["families_reconstructed"]) stats["production_docs_missing_linked_text"] += int(production_result["docs_missing_linked_text"]) stats["production_docs_missing_linked_images"] += int(production_result["docs_missing_linked_images"]) stats["production_docs_missing_linked_natives"] += int(production_result["docs_missing_linked_natives"]) for failure in list(production_result.get("failures", [])): control_number = normalize_whitespace(str(failure.get("control_number") or "")) failure_entry: dict[str, str] = { "rel_path": ( production_logical_rel_path(production_rel_root, control_number).as_posix() if control_number else production_rel_root ), "production_rel_root": production_rel_root, "error": str(failure.get("error") or ""), } if control_number: failure_entry["control_number"] = control_number failures.append(failure_entry) stats["failed"] += len(list(production_result.get("failures", []))) except Exception as exc: rollback_open_transaction(connection) stats["failed"] += 1 failures.append( { "rel_path": production_rel_root, "production_rel_root": production_rel_root, "error": f"{type(exc).__name__}: {exc}", } ) finally: production_ms += (time.perf_counter() - source_started) * 1000.0 else: skipped_production_roots = [str(signature["rel_root"]) for signature in production_signatures] warnings = [ f"Detected processed production root at {signature['rel_root']}; use ingest-production instead." for signature in production_signatures ] return { "current_ingestion_batch": current_ingestion_batch, "slack_day_documents_missing": slack_day_documents_missing, "ingested_production_roots": ingested_production_roots, "skipped_production_roots": skipped_production_roots, "warnings": warnings, "gmail_ms": gmail_ms, "slack_ms": slack_ms, "production_ms": production_ms, } def load_loose_file_commit_state( connection: sqlite3.Connection, scanned_rel_paths: set[str], gmail_owned_rel_paths: set[str], pst_export_owned_rel_paths: set[str], scan_scope: dict[str, object] | None = None, ) -> tuple[dict[str, sqlite3.Row], dict[str, list[sqlite3.Row]]]: existing_occurrence_rows = connection.execute( """ SELECT * FROM document_occurrences WHERE parent_occurrence_id IS NULL AND source_kind = ? AND lifecycle_status != 'deleted' """, (FILESYSTEM_SOURCE_KIND,), ).fetchall() existing_occurrence_rows = [ row for row in existing_occurrence_rows if str(row["rel_path"]) not in gmail_owned_rel_paths and str(row["rel_path"]) not in pst_export_owned_rel_paths ] existing_by_rel = {str(row["rel_path"]): row for row in existing_occurrence_rows} unseen_existing_by_hash: dict[str, list[sqlite3.Row]] = defaultdict(list) for row in existing_occurrence_rows: if ( ingest_scan_scope_contains_rel_path(scan_scope, str(row["rel_path"])) and str(row["rel_path"]) not in scanned_rel_paths and row["file_hash"] ): unseen_existing_by_hash[row["file_hash"]].append(row) return existing_by_rel, unseen_existing_by_hash def ingest_prepare_worker_count() -> int: default_workers = min(8, os.cpu_count() or 4) raw_value = os.environ.get("RETRIEVER_INGEST_WORKERS") if raw_value is None: return default_workers raw_value = raw_value.strip() if not raw_value: return default_workers try: return max(1, int(raw_value)) except ValueError: return default_workers def ingest_container_prepare_worker_count() -> int: default_workers = min(8, os.cpu_count() or 4) raw_value = os.environ.get("RETRIEVER_INGEST_CONTAINER_WORKERS") if raw_value is None: return default_workers raw_value = raw_value.strip() if not raw_value: return default_workers try: return max(1, int(raw_value)) except ValueError: return default_workers def ingest_prepare_queue_capacity(prepare_workers: int) -> int: return max(2, prepare_workers * 2) def ingest_prepare_queue_max_bytes() -> int: raw_value = os.environ.get("RETRIEVER_INGEST_PREPARED_QUEUE_BYTES") default_value = 512 * 1024 * 1024 if raw_value is None: return default_value raw_value = raw_value.strip() if not raw_value: return default_value try: return max(1, int(raw_value)) except ValueError: return default_value def ingest_prepare_spill_threshold_bytes() -> int: raw_value = os.environ.get("RETRIEVER_INGEST_PREPARE_SPILL_BYTES") default_value = 32 * 1024 * 1024 if raw_value is None: return default_value raw_value = raw_value.strip() if not raw_value: return default_value try: return max(1, int(raw_value)) except ValueError: return default_value def serialize_prepared_ingest_item(prepared_item: dict[str, object]) -> bytes: return pickle.dumps(prepared_item, protocol=pickle.HIGHEST_PROTOCOL) def stage_prepared_ingest_item( spill_dir: Path, prepared_index: int, serialized_payload: bytes, ) -> Path: spill_dir.mkdir(parents=True, exist_ok=True) spill_path = spill_dir / f"{prepared_index:08d}.prepared" spill_path.write_bytes(serialized_payload) return spill_path def hydrate_staged_prepared_ingest_item(spill_path: Path) -> dict[str, object]: try: return pickle.loads(spill_path.read_bytes()) finally: remove_file_if_exists(spill_path) def iter_staged_prepared_items( items: list[dict[str, object]] | Iterator[dict[str, object]], *, prepare_item, config_benchmark_name: str, queue_done_benchmark_name: str, spill_subdir_name: str, staging_root: Path | None = None, prepare_workers: int | None = None, ) -> Iterator[tuple[dict[str, object], float]]: effective_prepare_workers = prepare_workers if prepare_workers is not None else ingest_prepare_worker_count() max_prepared_items = ingest_prepare_queue_capacity(effective_prepare_workers) queue_max_bytes = ingest_prepare_queue_max_bytes() spill_threshold_bytes = ingest_prepare_spill_threshold_bytes() benchmark_mark( config_benchmark_name, prepare_workers=effective_prepare_workers, max_prepared_items=max_prepared_items, queue_max_bytes=queue_max_bytes, spill_threshold_bytes=spill_threshold_bytes, ) own_staging_dir = staging_root is None spill_dir = ( Path(tempfile.mkdtemp(prefix="retriever-ingest-prepared-")) if staging_root is None else Path(staging_root) / spill_subdir_name ) ready_entries_by_index: dict[int, dict[str, object]] = {} ready_bytes_in_memory = 0 peak_ready_bytes_in_memory = 0 spilled_items = 0 spilled_bytes = 0 try: with ThreadPoolExecutor(max_workers=effective_prepare_workers) as executor: pending_by_index: dict[int, object] = {} future_to_index: dict[object, int] = {} next_yield_index = 0 item_iterator = enumerate(items) input_exhausted = False def submit_available() -> None: nonlocal input_exhausted while ( not input_exhausted and len(pending_by_index) + len(ready_entries_by_index) < max_prepared_items ): try: prepared_index, item = next(item_iterator) except StopIteration: input_exhausted = True break future = executor.submit( prepare_item, item, ) pending_by_index[prepared_index] = future future_to_index[future] = prepared_index def record_completed_future(future) -> None: nonlocal ready_bytes_in_memory, peak_ready_bytes_in_memory, spilled_items, spilled_bytes prepared_index = future_to_index.pop(future) pending_by_index.pop(prepared_index, None) prepared_item = future.result() serialized_payload = serialize_prepared_ingest_item(prepared_item) serialized_size = len(serialized_payload) should_spill = ( serialized_size >= spill_threshold_bytes or ready_bytes_in_memory + serialized_size > queue_max_bytes ) if should_spill: spill_path = stage_prepared_ingest_item(spill_dir, prepared_index, serialized_payload) ready_entries_by_index[prepared_index] = { "storage": "spill", "spill_path": spill_path, "serialized_size": serialized_size, } spilled_items += 1 spilled_bytes += serialized_size return ready_entries_by_index[prepared_index] = { "storage": "memory", "prepared_item": prepared_item, "serialized_size": serialized_size, } ready_bytes_in_memory += serialized_size peak_ready_bytes_in_memory = max(peak_ready_bytes_in_memory, ready_bytes_in_memory) submit_available() while next_yield_index in ready_entries_by_index or pending_by_index or not input_exhausted: wait_started = time.perf_counter() while next_yield_index not in ready_entries_by_index: if not pending_by_index and input_exhausted: return if not pending_by_index: submit_available() if not pending_by_index and input_exhausted: return if not pending_by_index: raise RetrieverError( f"Prepared ingest queue drained before index {next_yield_index} was ready." ) done, _ = wait(list(pending_by_index.values()), return_when=FIRST_COMPLETED) for future in done: record_completed_future(future) submit_available() wait_ms = (time.perf_counter() - wait_started) * 1000.0 entry = ready_entries_by_index.pop(next_yield_index) if entry["storage"] == "spill": prepared_item = hydrate_staged_prepared_ingest_item(Path(entry["spill_path"])) else: prepared_item = dict(entry["prepared_item"]) ready_bytes_in_memory = max(0, ready_bytes_in_memory - int(entry["serialized_size"])) yield prepared_item, wait_ms next_yield_index += 1 submit_available() finally: benchmark_mark( queue_done_benchmark_name, spilled_items=spilled_items, spilled_bytes=spilled_bytes, peak_ready_bytes=peak_ready_bytes_in_memory, spill_dir=str(spill_dir), ) if own_staging_dir: remove_directory_tree(spill_dir) def prepare_loose_file_item(item: dict[str, object]) -> dict[str, object]: prepared_item = dict(item) prepare_started = time.perf_counter() try: extract_started = time.perf_counter() extracted_payload = extract_document(Path(item["path"]), include_attachments=True) prepared_item["prepare_extract_ms"] = ingest_v2_elapsed_ms(extract_started) attachments = list(extracted_payload.get("attachments", [])) extracted_payload = dict(extracted_payload) extracted_payload.pop("attachments", None) chunk_started = time.perf_counter() prepared_chunks = extracted_search_chunks(extracted_payload) prepared_item["extracted_payload"] = extracted_payload prepared_item["attachments"] = attachments prepared_item["prepared_chunks"] = prepared_chunks prepared_item["prepare_chunk_ms"] = (time.perf_counter() - chunk_started) * 1000.0 prepared_item["prepare_error"] = None except Exception as exc: prepared_item["extracted_payload"] = None prepared_item["attachments"] = [] prepared_item["prepared_chunks"] = [] prepared_item["prepare_extract_ms"] = 0.0 prepared_item["prepare_chunk_ms"] = 0.0 prepared_item["prepare_error"] = f"{type(exc).__name__}: {exc}" prepared_item["prepare_ms"] = (time.perf_counter() - prepare_started) * 1000.0 return prepared_item def refresh_prepared_loose_file_item_if_stale( prepared_item: dict[str, object], ) -> tuple[dict[str, object], bool]: path = Path(prepared_item["path"]) current_file_size, current_file_mtime_ns = source_file_snapshot(path) if ( current_file_size == prepared_item.get("source_file_size") and current_file_mtime_ns == prepared_item.get("source_file_mtime_ns") ): return prepared_item, False benchmark_mark( "ingest_loose_file_freshness_fallback", rel_path=str(prepared_item["rel_path"]), planned_file_size=prepared_item.get("source_file_size"), current_file_size=current_file_size, planned_file_mtime_ns=prepared_item.get("source_file_mtime_ns"), current_file_mtime_ns=current_file_mtime_ns, ) refreshed_item = refresh_ingest_item_filesystem_facts(prepared_item) return prepare_loose_file_item(refreshed_item), True def iter_prepared_loose_file_items( loose_file_items: list[dict[str, object]], staging_root: Path | None = None, ) -> Iterator[tuple[dict[str, object], float]]: yield from iter_staged_prepared_items( loose_file_items, prepare_item=prepare_loose_file_item, config_benchmark_name="ingest_loose_prepare_config", queue_done_benchmark_name="ingest_loose_prepare_queue_done", spill_subdir_name="prepared-loose", staging_root=staging_root, ) def commit_prepared_loose_file( connection: sqlite3.Connection, paths: dict[str, Path], prepared_item: dict[str, object], existing_by_rel: dict[str, sqlite3.Row], unseen_existing_by_hash: dict[str, list[sqlite3.Row]], ensure_filesystem_dataset, current_ingestion_batch: int | None, before_transaction_commit=None, ) -> dict[str, object]: prepared_item, freshness_fallback = refresh_prepared_loose_file_item_if_stale(prepared_item) rel_path = str(prepared_item["rel_path"]) path = Path(prepared_item["path"]) file_type = str(prepared_item["file_type"]) file_hash = prepared_item.get("file_hash") existing_occurrence_row = existing_by_rel.get(rel_path) action = "new" if existing_occurrence_row is not None: existing_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (existing_occurrence_row["document_id"],), ).fetchone() if existing_row is None: raise RetrieverError(f"Occurrence {existing_occurrence_row['id']} points at a missing document.") extracted_for_skip_check = prepared_item.get("extracted_payload") if isinstance(extracted_for_skip_check, dict): extracted_for_skip_check = apply_manual_locks(existing_row, dict(extracted_for_skip_check)) if ( existing_occurrence_row["file_hash"] == file_hash and existing_occurrence_row["lifecycle_status"] == ACTIVE_OCCURRENCE_STATUS and document_row_has_seeded_text_revisions(existing_row) and extracted_payload_matches_document_row(existing_row, extracted_for_skip_check) and document_row_has_email_threading(connection, existing_row) ): filesystem_dataset_id, filesystem_dataset_source_id = ensure_filesystem_dataset() connection.execute("BEGIN") try: mark_seen_without_reingest( connection, existing_occurrence_row, dataset_id=filesystem_dataset_id, dataset_source_id=filesystem_dataset_source_id, ) result = { "action": "skipped", "current_ingestion_batch": current_ingestion_batch, "freshness_fallback": freshness_fallback, "document_id": int(existing_occurrence_row["document_id"]), "file_hash": file_hash, } if before_transaction_commit is not None: before_transaction_commit(connection, result) connection.commit() return result except Exception: connection.rollback() raise action = "updated" else: existing_row = None rename_candidates = unseen_existing_by_hash.get(str(file_hash)) or [] if rename_candidates: existing_occurrence_row = rename_candidates.pop(0) action = "renamed" prepare_error = prepared_item.get("prepare_error") if prepare_error: return { "action": "failed", "current_ingestion_batch": current_ingestion_batch, "error": str(prepare_error), "freshness_fallback": freshness_fallback, } if existing_occurrence_row is None and file_hash: exact_duplicate_document = get_document_by_dedupe_key( connection, basis="file_hash", key_value=str(file_hash), ) if exact_duplicate_document is not None: filesystem_dataset_id, filesystem_dataset_source_id = ensure_filesystem_dataset() connection.execute("BEGIN") try: now = utc_now() duplicate_occurrence_id = attach_occurrence_to_existing_document( connection, exact_duplicate_document, existing_occurrence_row=None, rel_path=rel_path, file_name=path.name, file_type=file_type, file_size=path.stat().st_size, file_hash=str(file_hash), source_kind=FILESYSTEM_SOURCE_KIND, source_rel_path=rel_path, source_item_id=None, source_folder_path=None, custodian=infer_source_custodian( source_kind=FILESYSTEM_SOURCE_KIND, source_rel_path=rel_path, ), occurrence_control_number=str(exact_duplicate_document["control_number"] or ""), ingested_at=now, last_seen_at=now, updated_at=now, ) clone_duplicate_family_child_occurrences( connection, paths, parent_document_id=int(exact_duplicate_document["id"]), parent_occurrence_id=duplicate_occurrence_id, parent_rel_path=rel_path, custodian=infer_source_custodian( source_kind=FILESYSTEM_SOURCE_KIND, source_rel_path=rel_path, ), ingested_at=now, last_seen_at=now, updated_at=now, ) ensure_dataset_document_membership( connection, dataset_id=filesystem_dataset_id, document_id=int(exact_duplicate_document["id"]), dataset_source_id=filesystem_dataset_source_id, ) result = { "action": "new", "current_ingestion_batch": current_ingestion_batch, "freshness_fallback": freshness_fallback, "document_id": int(exact_duplicate_document["id"]), "file_hash": file_hash, } if before_transaction_commit is not None: before_transaction_commit(connection, result) connection.commit() return result except Exception: connection.rollback() raise filesystem_dataset_id, filesystem_dataset_source_id = ensure_filesystem_dataset() connection.execute("BEGIN") try: existing_row = None reused_existing_occurrence_row = existing_occurrence_row superseded_document_id: int | None = None if existing_occurrence_row is not None: existing_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (existing_occurrence_row["document_id"],), ).fetchone() if existing_row is None: raise RetrieverError(f"Occurrence {existing_occurrence_row['id']} points at a missing document.") active_occurrence_rows = active_occurrence_rows_for_document(connection, int(existing_row["id"])) if ( action == "updated" and len(active_occurrence_rows) > 1 and existing_occurrence_row["file_hash"] != file_hash ): connection.execute( """ UPDATE document_occurrences SET lifecycle_status = 'superseded', updated_at = ? WHERE id = ? """, (utc_now(), existing_occurrence_row["id"]), ) superseded_document_id = int(existing_row["id"]) refresh_source_backed_dataset_memberships_for_document(connection, superseded_document_id) refresh_document_from_occurrences(connection, superseded_document_id) existing_row = None reused_existing_occurrence_row = None extracted = apply_manual_locks(existing_row, dict(prepared_item["extracted_payload"] or {})) attachments = list(prepared_item.get("attachments", [])) if existing_row is None: if current_ingestion_batch is None: current_ingestion_batch = allocate_ingestion_batch_number(connection) control_number_batch = current_ingestion_batch control_number_family_sequence = reserve_control_number_family_sequence(connection, control_number_batch) control_number = format_control_number(control_number_batch, control_number_family_sequence) control_number_attachment_sequence = None else: control_number_batch = int(existing_row["control_number_batch"]) control_number_family_sequence = int(existing_row["control_number_family_sequence"]) control_number = str(existing_row["control_number"]) control_number_attachment_sequence = existing_row["control_number_attachment_sequence"] cleanup_document_artifacts(paths, connection, existing_row) document_id = upsert_document_row( connection, rel_path, path, existing_row, extracted, existing_occurrence_row=reused_existing_occurrence_row, file_name=path.name, parent_document_id=None, control_number=control_number, dataset_id=filesystem_dataset_id, control_number_batch=control_number_batch, control_number_family_sequence=control_number_family_sequence, control_number_attachment_sequence=control_number_attachment_sequence, ) replace_document_email_threading_row( connection, document_id=document_id, email_threading=extracted.get("email_threading"), ) seed_source_text_revision_for_document( connection, paths, document_id=document_id, extracted=extracted, existing_row=existing_row, ) ensure_dataset_document_membership( connection, dataset_id=filesystem_dataset_id, document_id=document_id, dataset_source_id=filesystem_dataset_source_id, ) preview_rows = write_preview_artifacts(paths, rel_path, list(extracted.get("preview_artifacts", []))) replace_document_related_rows( connection, document_id, extracted | {"file_name": path.name}, list(prepared_item.get("prepared_chunks", [])), preview_rows, ) reconcile_attachment_documents( connection, paths, document_id, rel_path, control_number_batch, control_number_family_sequence, attachments, [(filesystem_dataset_id, filesystem_dataset_source_id)], ) if superseded_document_id is not None and superseded_document_id != document_id: refresh_source_backed_dataset_memberships_for_document(connection, superseded_document_id) refresh_document_from_occurrences(connection, superseded_document_id) result = { "action": action, "current_ingestion_batch": current_ingestion_batch, "freshness_fallback": freshness_fallback, "document_id": document_id, "file_hash": file_hash, } if action == "renamed" and existing_occurrence_row is not None: result["source_occurrence_id"] = int(existing_occurrence_row["id"]) result["source_document_id"] = int(existing_occurrence_row["document_id"]) if before_transaction_commit is not None: before_transaction_commit(connection, result) connection.commit() return result except Exception as exc: connection.rollback() return { "action": "failed", "current_ingestion_batch": current_ingestion_batch, "error": f"{type(exc).__name__}: {exc}", "freshness_fallback": freshness_fallback, } def finalize_ingest_postpass( connection: sqlite3.Connection, paths: dict[str, Path], allowed_types: set[str] | None, scanned_rel_paths: set[str], scanned_pst_source_rel_paths: set[str], scanned_mbox_source_rel_paths: set[str], slack_day_documents_missing: int, stats: dict[str, int], pst_export_owned_rel_paths: set[str] | None = None, scan_scope: dict[str, object] | None = None, ) -> int: if pst_export_owned_rel_paths: connection.execute("BEGIN") try: retire_standalone_filesystem_documents_by_rel_paths( connection, paths, rel_paths=pst_export_owned_rel_paths, ) connection.commit() except Exception: connection.rollback() raise filesystem_missing = mark_missing_documents(connection, scanned_rel_paths, scan_scope=scan_scope) pst_sources_missing = 0 pst_documents_missing = 0 mbox_sources_missing = 0 mbox_documents_missing = 0 if allowed_types is None or PST_SOURCE_KIND in allowed_types: pst_sources_missing, pst_documents_missing = mark_missing_pst_documents( connection, scanned_pst_source_rel_paths, scan_scope=scan_scope, ) if allowed_types is None or MBOX_SOURCE_KIND in allowed_types: mbox_sources_missing, mbox_documents_missing = mark_missing_mbox_documents( connection, scanned_mbox_source_rel_paths, scan_scope=scan_scope, ) stats["pst_sources_missing"] = pst_sources_missing stats["pst_documents_missing"] = pst_documents_missing stats["mbox_sources_missing"] = mbox_sources_missing stats["mbox_documents_missing"] = mbox_documents_missing stats["slack_documents_missing"] = slack_day_documents_missing stats["missing"] = filesystem_missing + pst_sources_missing + mbox_sources_missing + slack_day_documents_missing connection.execute("BEGIN") try: conversation_assignment = assign_supported_conversations(connection) refresh_conversation_previews(connection, paths) connection.commit() except Exception: connection.rollback() raise stats["email_conversations"] = int(conversation_assignment["email_conversations"]) stats["email_documents_reassigned"] = int(conversation_assignment["email_documents_reassigned"]) stats["email_child_documents_updated"] = int(conversation_assignment["email_child_documents_updated"]) stats["pst_chat_conversations"] = int(conversation_assignment["pst_chat_conversations"]) stats["pst_chat_documents_reassigned"] = int(conversation_assignment["pst_chat_documents_reassigned"]) stats["pst_chat_child_documents_updated"] = int(conversation_assignment["pst_chat_child_documents_updated"]) connection.execute("BEGIN") try: pruned_unused_filesystem_dataset = prune_unused_filesystem_dataset(connection) connection.commit() except Exception: connection.rollback() raise return int(pruned_unused_filesystem_dataset) def ingest_production(root: Path, production_root: Path | str) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) ensure_layout(paths) total_started = time.perf_counter() benchmark_mark("ingest_production_begin") with workspace_ingest_session(paths, command_name="ingest-production") as ingest_session: connection = connect_db(paths["db_path"]) try: setup_started = time.perf_counter() apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="ingest-production") reconcile_custom_fields_registry(connection, repair=True) resolved_production_root = resolve_production_root_argument(root, production_root) benchmark_mark( "ingest_production_setup_done", setup_ms=round((time.perf_counter() - setup_started) * 1000.0, 3), production_root=str(resolved_production_root), ) production_started = time.perf_counter() result = ingest_resolved_production_root( connection, paths, root, resolved_production_root, staging_root=Path(ingest_session["tmp_dir"]), ) session_warnings = list(ingest_session.get("warnings") or []) if session_warnings: result = {**result, "warnings": [*list(result.get("warnings", [])), *session_warnings]} benchmark_mark( "ingest_production_done", production_ms=round((time.perf_counter() - production_started) * 1000.0, 3), total_ms=round((time.perf_counter() - total_started) * 1000.0, 3), created=int(result.get("created") or 0), updated=int(result.get("updated") or 0), failed=len(list(result.get("failures", []))), ) return result except Exception as exc: benchmark_mark( "ingest_production_failed", total_ms=round((time.perf_counter() - total_started) * 1000.0, 3), error=f"{type(exc).__name__}: {exc}", ) raise finally: connection.close() def ingest_v2_phase_cursor_payload( connection: sqlite3.Connection, *, run_id: str, phase: str, cursor_key: str, ) -> dict[str, object]: row = connection.execute( """ SELECT cursor_json FROM ingest_phase_cursors WHERE run_id = ? AND phase = ? AND cursor_key = ? """, (run_id, phase, cursor_key), ).fetchone() if row is None: return {} payload = decode_json_text(row["cursor_json"], default={}) or {} return payload if isinstance(payload, dict) else {} def ingest_v2_compat_failures( connection: sqlite3.Connection, *, run_id: str, plan_cursor: dict[str, object], ) -> list[dict[str, str]]: failures: list[dict[str, str]] = [] seen_entries: set[tuple[str, str]] = set() def append_failure(rel_path: str, error: str) -> None: normalized_rel_path = normalize_whitespace(rel_path) normalized_error = normalize_whitespace(error) if not normalized_rel_path or not normalized_error: return key = (normalized_rel_path, normalized_error) if key in seen_entries: return seen_entries.add(key) failures.append({"rel_path": normalized_rel_path, "error": normalized_error}) for failure_key, rel_path_key in ( ("slack_failures", "slack_rel_root"), ("mbox_failures", "source_rel_path"), ("pst_failures", "source_rel_path"), ("production_failures", "production_rel_root"), ): for item in list(plan_cursor.get(failure_key) or []): if not isinstance(item, dict): continue append_failure( str(item.get(rel_path_key) or ""), str(item.get("error") or ""), ) failed_rows = connection.execute( """ SELECT rel_path, source_key, unit_type, last_error FROM ingest_work_items WHERE run_id = ? AND status = 'failed' ORDER BY id ASC """, (run_id,), ).fetchall() for row in failed_rows: append_failure( str(row["rel_path"] or row["source_key"] or row["unit_type"] or ""), str(row["last_error"] or ""), ) return failures def ingest_v2_compat_action_counts( connection: sqlite3.Connection, *, run_id: str, plan_cursor: dict[str, object], ) -> dict[str, int]: counts = { "new": 0, "updated": 0, "renamed": 0, "skipped": int(plan_cursor.get("skipped_unchanged_loose_files") or 0), } rows = connection.execute( """ SELECT unit_type, payload_json, artifact_manifest_json FROM ingest_work_items WHERE run_id = ? AND status = 'committed' ORDER BY id ASC """, (run_id,), ).fetchall() for row in rows: unit_type = str(row["unit_type"] or "") payload = decode_json_text(row["payload_json"], default={}) or {} payload = payload if isinstance(payload, dict) else {} manifest = decode_json_text(row["artifact_manifest_json"], default={}) or {} manifest = manifest if isinstance(manifest, dict) else {} if unit_type in {"mbox_source_finalizer", "pst_source_finalizer"}: source_action = normalize_whitespace( str(payload.get("source_action") or manifest.get("source_action") or "") ) if source_action in counts: counts[source_action] += 1 continue if unit_type in { "mbox_message", "pst_message", "conversation_preview", "production_preview_batch", "production_row", }: continue if unit_type in {"slack_conversation", "slack_document"}: counts["new"] += int(manifest.get("new") or 0) counts["updated"] += int(manifest.get("updated") or 0) continue commit_action = normalize_whitespace(str(manifest.get("commit_action") or "")) if commit_action in counts: counts[commit_action] += 1 return counts def ingest_v2_compat_summary( connection: sqlite3.Connection, root: Path, *, run_id: str, ) -> dict[str, object]: run_row = require_ingest_v2_run_row(connection, run_id) scope_payload = decode_json_text(run_row["scope_json"], default={}) or {} scope = scope_payload if isinstance(scope_payload, dict) else {} plan_cursor = ingest_v2_phase_cursor_payload( connection, run_id=run_id, phase="planning", cursor_key="loose_file_scan", ) commit_cursor = ingest_v2_phase_cursor_payload( connection, run_id=run_id, phase="committing", cursor_key="loose_file_commit", ) finalize_cursor = ingest_v2_phase_cursor_payload( connection, run_id=run_id, phase="finalizing", cursor_key="loose_file_finalize", ) compat_actions = ingest_v2_compat_action_counts(connection, run_id=run_id, plan_cursor=plan_cursor) special_source_counts = dict(plan_cursor.get("special_source_counts") or {}) slack_stats = dict(commit_cursor.get("slack_stats") or {}) mbox_stats = dict(commit_cursor.get("mbox_stats") or {}) pst_stats = dict(commit_cursor.get("pst_stats") or {}) production_commit_stats = dict(commit_cursor.get("production_stats") or {}) production_finalize_stats = dict(finalize_cursor.get("production_stats") or {}) conversation_assignment = dict(finalize_cursor.get("conversation_assignment") or {}) planned_slack_export_roots = list(plan_cursor.get("planned_slack_export_roots") or []) planned_gmail_mbox_sources = int(plan_cursor.get("planned_gmail_mbox_sources") or 0) planned_mbox_sources = list(plan_cursor.get("planned_mbox_sources") or []) planned_pst_sources = list(plan_cursor.get("planned_pst_sources") or []) planned_production_roots = list(plan_cursor.get("planned_production_roots") or []) skipped_production_roots = list(plan_cursor.get("skipped_production_roots") or []) failures = ingest_v2_compat_failures(connection, run_id=run_id, plan_cursor=plan_cursor) workspace_inventory = document_inventory_counts(connection) warnings = [ f"Detected processed production root at {production_rel_root}; use ingest-production instead." for production_rel_root in skipped_production_roots if normalize_whitespace(str(production_rel_root or "")) ] scanned_files = ( int(plan_cursor.get("planned_loose_files") or 0) + int(plan_cursor.get("skipped_unchanged_loose_files") or 0) + len(planned_mbox_sources) + len(planned_pst_sources) + int(plan_cursor.get("planned_slack_day_files_scanned") or 0) + planned_gmail_mbox_sources ) result = default_ingest_stats( int(special_source_counts.get("slack_export_roots") or len(planned_slack_export_roots)), int(special_source_counts.get("gmail_export_roots") or 0), ) result.update( { "new": int(compat_actions.get("new") or 0) + int(production_commit_stats.get("created") or 0), "updated": int(compat_actions.get("updated") or 0) + int(production_commit_stats.get("updated") or 0), "renamed": int(compat_actions.get("renamed") or 0), "skipped": int(compat_actions.get("skipped") or 0), "failed": len(failures), "missing": ( int(finalize_cursor.get("filesystem_missing") or 0) + int(finalize_cursor.get("pst_sources_missing") or 0) + int(finalize_cursor.get("mbox_sources_missing") or 0) + int(finalize_cursor.get("slack_documents_missing") or 0) ), "pst_sources_skipped": int(pst_stats.get("pst_sources_skipped") or 0), "pst_messages_created": int(pst_stats.get("pst_messages_created") or 0), "pst_messages_updated": int(pst_stats.get("pst_messages_updated") or 0), "pst_messages_deleted": int(pst_stats.get("pst_messages_deleted") or 0), "pst_sources_missing": int(finalize_cursor.get("pst_sources_missing") or 0), "pst_documents_missing": int(finalize_cursor.get("pst_documents_missing") or 0), "mbox_sources_skipped": int(mbox_stats.get("mbox_sources_skipped") or 0), "mbox_messages_created": int(mbox_stats.get("mbox_messages_created") or 0), "mbox_messages_updated": int(mbox_stats.get("mbox_messages_updated") or 0), "mbox_messages_deleted": int(mbox_stats.get("mbox_messages_deleted") or 0), "mbox_sources_missing": int(finalize_cursor.get("mbox_sources_missing") or 0), "mbox_documents_missing": int(finalize_cursor.get("mbox_documents_missing") or 0), "gmail_documents_scanned": planned_gmail_mbox_sources, "slack_day_documents_scanned": int(plan_cursor.get("planned_slack_day_files_scanned") or 0), "slack_documents_created": int(slack_stats.get("slack_documents_created") or 0), "slack_documents_updated": int(slack_stats.get("slack_documents_updated") or 0), "slack_documents_missing": int(finalize_cursor.get("slack_documents_missing") or 0), "slack_conversations": int(slack_stats.get("slack_conversations") or 0), "email_conversations": int(conversation_assignment.get("email_conversations") or 0), "email_documents_reassigned": int(conversation_assignment.get("email_documents_reassigned") or 0), "email_child_documents_updated": int(conversation_assignment.get("email_child_documents_updated") or 0), "pst_chat_conversations": int(conversation_assignment.get("pst_chat_conversations") or 0), "pst_chat_documents_reassigned": int(conversation_assignment.get("pst_chat_documents_reassigned") or 0), "pst_chat_child_documents_updated": int( conversation_assignment.get("pst_chat_child_documents_updated") or 0 ), "production_documents_created": int(production_commit_stats.get("created") or 0), "production_documents_updated": int(production_commit_stats.get("updated") or 0), "production_documents_unchanged": int(production_commit_stats.get("unchanged") or 0), "production_documents_retired": int(production_finalize_stats.get("retired") or 0), "production_families_reconstructed": int( production_finalize_stats.get("families_reconstructed") or 0 ), "production_docs_missing_linked_text": int( plan_cursor.get("production_docs_missing_linked_text") or 0 ), "production_docs_missing_linked_images": int( plan_cursor.get("production_docs_missing_linked_images") or 0 ), "production_docs_missing_linked_natives": int( plan_cursor.get("production_docs_missing_linked_natives") or 0 ), } ) result["failures"] = failures result["scanned"] = scanned_files result["scanned_files"] = scanned_files if not bool(scope.get("is_full_workspace")): result["scan_paths"] = list(scope.get("scan_paths") or []) result["pruned_unused_filesystem_dataset"] = int( bool(finalize_cursor.get("pruned_unused_filesystem_dataset") or False) ) result["ingested_production_roots"] = list(finalize_cursor.get("production_finalized_roots") or planned_production_roots) result["skipped_production_roots"] = skipped_production_roots if warnings: result["warnings"] = warnings result["workspace_parent_documents"] = workspace_inventory["parent_documents"] result["workspace_missing_parent_documents"] = workspace_inventory["missing_parent_documents"] result["workspace_attachment_children"] = workspace_inventory["attachment_children"] result["workspace_documents_total"] = workspace_inventory["documents_total"] result["run_id"] = run_id result["status"] = str(run_row["status"] or "") return result def ingest( root: Path, recursive: bool, raw_file_types: str | None, raw_paths: list[str] | None = None, ) -> dict[str, object]: set_active_workspace_root(root) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="ingest") reconcile_custom_fields_registry(connection, repair=True) finally: connection.close() facade_payload = ingest_v2_facade( root, recursive=recursive, raw_file_types=raw_file_types, raw_paths=raw_paths, run_to_completion=True, skip_unchanged_loose_files=False, ) run_id = normalize_whitespace(str(facade_payload.get("run_id") or "")) if not run_id: raise RetrieverError("Ingest completed without a resumable run id.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) return ingest_v2_compat_summary(connection, root, run_id=run_id) finally: connection.close() def value_from_type(field_type: str, value: str | None) -> object: if value is None: return None if field_type == "date": normalized = normalize_date_field_value(value) if normalized is None: raise RetrieverError(f"Expected ISO date value, got {value!r}") return normalized if field_type == "integer": try: return int(value) except ValueError as exc: raise RetrieverError(f"Expected integer value, got {value!r}") from exc if field_type == "real": try: return float(value) except ValueError as exc: raise RetrieverError(f"Expected real value, got {value!r}") from exc if field_type == "boolean": lowered = value.strip().lower() truthy = {"1", "true", "yes", "y", "on"} falsy = {"0", "false", "no", "n", "off"} if lowered in truthy: return 1 if lowered in falsy: return 0 raise RetrieverError(f"Expected boolean value, got {value!r}") return value def resolve_field_definition(connection: sqlite3.Connection, field_name: str) -> dict[str, str]: field_name = FIELD_NAME_ALIASES.get(field_name, field_name) if field_name in BUILTIN_FIELD_TYPES: return { "field_name": field_name, "field_type": BUILTIN_FIELD_TYPES[field_name], "source": "builtin", "displayable": "true", } if field_name in VIRTUAL_FILTER_FIELD_TYPES: return { "field_name": field_name, "field_type": VIRTUAL_FILTER_FIELD_TYPES[field_name], "source": "virtual", "displayable": "true" if field_name in DISPLAYABLE_VIRTUAL_FIELDS else "false", } row = connection.execute( """ SELECT field_name, field_type FROM custom_fields_registry WHERE field_name = ? """, (field_name,), ).fetchone() columns = table_columns(connection, "documents") if row is not None and row["field_name"] in columns: return { "field_name": row["field_name"], "field_type": row["field_type"], "source": "custom", "displayable": "true", } if field_name in columns: sqlite_type = next( (info["type"] for info in table_info(connection, "documents") if info["name"] == field_name), "", ) return { "field_name": field_name, "field_type": infer_registry_field_type(sqlite_type), "source": "column", "displayable": "true", } raise RetrieverError(f"Unknown field: {field_name}") def lock_field(connection: sqlite3.Connection, document_id: int, field_name: str) -> list[str]: row = connection.execute( f"SELECT {MANUAL_FIELD_LOCKS_COLUMN} FROM documents WHERE id = ?", (document_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown document id: {document_id}") locks = normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN]) if field_name not in locks: locks.append(field_name) connection.execute( f"UPDATE documents SET {MANUAL_FIELD_LOCKS_COLUMN} = ?, updated_at = ? WHERE id = ?", (json.dumps(locks), utc_now(), document_id), ) return locks def row_to_plain_dict(row: sqlite3.Row | None) -> dict[str, object] | None: if row is None: return None return {key: row[key] for key in row.keys()} def normalize_merge_field_value(value: object) -> object: if value is None: return None if isinstance(value, str): normalized = normalize_whitespace(value) return normalized or None return value def serialize_merge_field_value(value: object) -> object: normalized = normalize_merge_field_value(value) if normalized is None: return None if isinstance(normalized, (str, int, float, bool)): return normalized return json.dumps(normalized, ensure_ascii=True, sort_keys=True) def merge_field_value_key(value: object) -> str | None: normalized = normalize_merge_field_value(value) if normalized is None: return None return json.dumps(normalized, ensure_ascii=True, sort_keys=True) def reconcile_custom_field_names(connection: sqlite3.Connection) -> list[str]: document_columns = table_columns(connection, "documents") rows = connection.execute( """ SELECT field_name FROM custom_fields_registry ORDER BY field_name ASC """ ).fetchall() return [str(row["field_name"]) for row in rows if str(row["field_name"]) in document_columns] def document_has_non_deleted_children(connection: sqlite3.Connection, document_id: int) -> list[int]: rows = connection.execute( """ SELECT id FROM documents WHERE parent_document_id = ? AND lifecycle_status != 'deleted' ORDER BY id ASC """, (document_id,), ).fetchall() return [int(row["id"]) for row in rows] def document_text_length(connection: sqlite3.Connection, document_id: int) -> int: row = connection.execute( """ SELECT COALESCE(MAX(char_end), 0) AS text_length FROM document_chunks WHERE document_id = ? """, (document_id,), ).fetchone() if row is None: return 0 return int(row["text_length"] or 0) def document_active_occurrence_count(connection: sqlite3.Connection, document_id: int) -> int: row = connection.execute( """ SELECT COUNT(*) AS occurrence_count FROM document_occurrences WHERE document_id = ? AND lifecycle_status = ? """, (document_id, ACTIVE_OCCURRENCE_STATUS), ).fetchone() if row is None: return 0 return int(row["occurrence_count"] or 0) def document_earliest_active_occurrence_ingested_at( connection: sqlite3.Connection, document_id: int, ) -> datetime: row = connection.execute( """ SELECT MIN(ingested_at) AS earliest_ingested_at FROM document_occurrences WHERE document_id = ? AND lifecycle_status = ? """, (document_id, ACTIVE_OCCURRENCE_STATUS), ).fetchone() parsed = parse_utc_timestamp(row["earliest_ingested_at"]) if row is not None else None return parsed or datetime.max.replace(tzinfo=timezone.utc) def document_canonical_field_count(row: sqlite3.Row | dict[str, object]) -> int: field_names = [ "author", "content_type", "date_created", "date_modified", "page_count", "participants", "recipients", "subject", "title", ] return sum(1 for field_name in field_names if normalize_merge_field_value(row[field_name]) is not None) def collect_distinct_document_field_values( document_rows: list[sqlite3.Row], field_name: str, ) -> list[dict[str, object]]: distinct_values: dict[str, dict[str, object]] = {} for row in document_rows: value = normalize_merge_field_value(row[field_name]) if value is None: continue value_key = merge_field_value_key(value) assert value_key is not None entry = distinct_values.setdefault( value_key, { "value": value, "document_ids": [], "locked_document_ids": [], }, ) entry["document_ids"].append(int(row["id"])) if field_name in normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN]): entry["locked_document_ids"].append(int(row["id"])) return list(distinct_values.values()) def choose_reconcile_survivor( connection: sqlite3.Connection, document_rows: list[sqlite3.Row], ) -> sqlite3.Row: ranked_rows = sorted( document_rows, key=lambda row: ( text_status_priority(row["text_status"]), -document_canonical_field_count(row), -document_text_length(connection, int(row["id"])), -document_active_occurrence_count(connection, int(row["id"])), document_earliest_active_occurrence_ingested_at(connection, int(row["id"])), int(row["id"]), ), ) return ranked_rows[0] def document_artifact_counts(connection: sqlite3.Connection, document_id: int) -> dict[str, int]: counts: dict[str, int] = {} for table_name in ( "document_chunks", "chunks_fts", "document_previews", "document_source_parts", "documents_fts", "text_revisions", ): row = connection.execute( f"SELECT COUNT(*) AS row_count FROM {table_name} WHERE document_id = ?", (document_id,), ).fetchone() counts[table_name] = int(row["row_count"] or 0) if row is not None else 0 return counts def document_merge_snapshot(connection: sqlite3.Connection, document_id: int) -> dict[str, object]: document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (document_id,), ).fetchone() occurrence_rows = connection.execute( """ SELECT * FROM document_occurrences WHERE document_id = ? ORDER BY id ASC """, (document_id,), ).fetchall() dataset_rows = connection.execute( """ SELECT * FROM dataset_documents WHERE document_id = ? ORDER BY dataset_id ASC, COALESCE(dataset_source_id, 0) ASC, id ASC """, (document_id,), ).fetchall() alias_rows = connection.execute( """ SELECT * FROM document_control_number_aliases WHERE document_id = ? ORDER BY id ASC """, (document_id,), ).fetchall() text_revision_rows = connection.execute( """ SELECT * FROM text_revisions WHERE document_id = ? ORDER BY id ASC """, (document_id,), ).fetchall() return { "schema_version": SCHEMA_VERSION, "document": row_to_plain_dict(document_row), "occurrences": [row_to_plain_dict(row) for row in occurrence_rows], "dataset_memberships": [row_to_plain_dict(row) for row in dataset_rows], "control_number_aliases": [row_to_plain_dict(row) for row in alias_rows], "text_revisions": [row_to_plain_dict(row) for row in text_revision_rows], } def insert_document_merge_event( connection: sqlite3.Connection, *, survivor_document_id: int, loser_document_id: int, merge_basis: str, pre_merge_survivor_snapshot: dict[str, object], pre_merge_loser_snapshot: dict[str, object], artifact_counts: dict[str, object], ) -> int: now = utc_now() connection.execute( """ INSERT INTO document_merge_events ( survivor_document_id, loser_document_id, merge_basis, actor, schema_version, pre_merge_survivor_json, pre_merge_loser_json, artifact_counts_json, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( survivor_document_id, loser_document_id, merge_basis, "reconcile-duplicates", SCHEMA_VERSION, json.dumps(pre_merge_survivor_snapshot, ensure_ascii=True, sort_keys=True), json.dumps(pre_merge_loser_snapshot, ensure_ascii=True, sort_keys=True), json.dumps(artifact_counts, ensure_ascii=True, sort_keys=True), now, ), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) def insert_document_field_conflicts( connection: sqlite3.Connection, *, merge_event_id: int, survivor_document_id: int, final_survivor_row: sqlite3.Row, pre_merge_survivor_row: sqlite3.Row, pre_merge_loser_row: sqlite3.Row, custom_field_names: list[str], ) -> None: tracked_fields = sorted(set(custom_field_names) | set(EDITABLE_BUILTIN_FIELDS)) rows_to_insert: list[tuple[int, int, str, object, object, str, str]] = [] now = utc_now() for field_name in tracked_fields: pre_survivor_value = normalize_merge_field_value(pre_merge_survivor_row[field_name]) loser_value = normalize_merge_field_value(pre_merge_loser_row[field_name]) if pre_survivor_value == loser_value: continue final_value = normalize_merge_field_value(final_survivor_row[field_name]) if final_value == loser_value and loser_value is not None: resolution = "adopted_loser" elif final_value == pre_survivor_value: resolution = "kept_survivor" else: resolution = "recomputed" rows_to_insert.append( ( merge_event_id, survivor_document_id, field_name, serialize_merge_field_value(pre_survivor_value), serialize_merge_field_value(loser_value), resolution, now, ) ) if rows_to_insert: connection.executemany( """ INSERT INTO document_field_conflicts ( merge_event_id, document_id, field_name, survivor_value, loser_value, resolution, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, rows_to_insert, ) def tombstone_rel_path_for_merged_document(document_id: int) -> str: return (Path(".retriever") / "merged" / f"{document_id}.merged").as_posix() def active_child_rows_for_parent_document( connection: sqlite3.Connection, parent_document_id: int, ) -> list[sqlite3.Row]: return connection.execute( """ SELECT * FROM documents WHERE parent_document_id = ? AND lifecycle_status = 'active' ORDER BY id ASC """, (parent_document_id,), ).fetchall() def family_child_signature_payload(row: sqlite3.Row) -> dict[str, object]: canonical_kind = normalize_whitespace(str(row["canonical_kind"] or "")).lower() return { "child_document_kind": normalize_whitespace(str(row["child_document_kind"] or "")).lower() or None, "file_name": normalize_whitespace(str(row["file_name"] or "")).lower() or None, "content_hash": normalize_whitespace(str(row["content_hash"] or "")) or None, "file_hash": normalize_whitespace(str(row["file_hash"] or "")) or None, "content_type": normalize_whitespace(str(row["content_type"] or "")) or None, "canonical_kind": canonical_kind if canonical_kind not in {"", "unknown"} else None, } def family_child_signature_key(row: sqlite3.Row) -> str: return sha256_json_value(family_child_signature_payload(row)) def document_family_descriptor( connection: sqlite3.Connection, document_row: sqlite3.Row, ) -> dict[str, object]: child_rows = active_child_rows_for_parent_document(connection, int(document_row["id"])) ordered_children = sorted( child_rows, key=lambda row: ( family_child_signature_key(row), int(row["id"]), ), ) child_slot_keys = [family_child_signature_key(row) for row in ordered_children] return { "document_id": int(document_row["id"]), "fingerprint": sha256_json_value({"child_slots": child_slot_keys}), "child_rows": ordered_children, "child_slot_keys": child_slot_keys, } def survivor_selection_payload( connection: sqlite3.Connection, survivor_row: sqlite3.Row, *, rule: str, ) -> dict[str, object]: survivor_id = int(survivor_row["id"]) return { "rule": rule, "text_status": survivor_row["text_status"], "canonical_field_count": document_canonical_field_count(survivor_row), "text_length": document_text_length(connection, survivor_id), "active_occurrence_count": document_active_occurrence_count(connection, survivor_id), "earliest_active_ingested_at": format_utc_timestamp( document_earliest_active_occurrence_ingested_at(connection, survivor_id) ), } def evaluate_document_merge_group( connection: sqlite3.Connection, document_rows: list[sqlite3.Row], *, custom_fields: list[str], forced_survivor_document_id: int | None = None, selection_rule: str | None = None, ) -> dict[str, object]: if forced_survivor_document_id is None: survivor_row = choose_reconcile_survivor(connection, document_rows) resolved_selection_rule = selection_rule or "text_status>field_count>text_length>occurrence_count>earliest_ingested_at>document_id" else: survivor_row = next( (row for row in document_rows if int(row["id"]) == int(forced_survivor_document_id)), None, ) if survivor_row is None: raise RetrieverError(f"Forced survivor document id {forced_survivor_document_id} is not present in the merge group.") resolved_selection_rule = selection_rule or "family_slot_survivor" survivor_id = int(survivor_row["id"]) loser_document_ids = [int(row["id"]) for row in document_rows if int(row["id"]) != survivor_id] survivor_locks = set(normalize_string_list(survivor_row[MANUAL_FIELD_LOCKS_COLUMN])) blockers: list[dict[str, object]] = [] machine_field_conflicts: list[dict[str, object]] = [] custom_field_updates: dict[str, object] = {} builtin_field_updates: dict[str, object] = {} fields_to_lock: set[str] = set(survivor_locks) non_unknown_kinds = sorted( { normalize_whitespace(str(row["canonical_kind"] or "")).lower() for row in document_rows if normalize_whitespace(str(row["canonical_kind"] or "")).lower() not in {"", "unknown"} } ) if len(non_unknown_kinds) > 1: blockers.append( { "type": "canonical_kind_conflict", "canonical_kinds": non_unknown_kinds, "document_ids": [int(row["id"]) for row in document_rows], } ) for field_name in custom_fields: distinct_values = collect_distinct_document_field_values(document_rows, field_name) if len(distinct_values) > 1: blockers.append( { "type": "custom_field_conflict", "field_name": field_name, "values": [ { "value": serialize_merge_field_value(item["value"]), "document_ids": item["document_ids"], } for item in distinct_values ], } ) continue if not distinct_values: continue chosen_value = distinct_values[0]["value"] if normalize_merge_field_value(survivor_row[field_name]) is None: custom_field_updates[field_name] = chosen_value if field_name in survivor_locks or distinct_values[0]["locked_document_ids"]: fields_to_lock.add(field_name) for field_name in sorted(EDITABLE_BUILTIN_FIELDS): distinct_values = collect_distinct_document_field_values(document_rows, field_name) if len(distinct_values) > 1 and any(item["locked_document_ids"] for item in distinct_values): blockers.append( { "type": "locked_builtin_conflict", "field_name": field_name, "values": [ { "value": serialize_merge_field_value(item["value"]), "document_ids": item["document_ids"], "locked_document_ids": item["locked_document_ids"], } for item in distinct_values ], } ) continue if len(distinct_values) > 1: machine_field_conflicts.append( { "field_name": field_name, "values": [ { "value": serialize_merge_field_value(item["value"]), "document_ids": item["document_ids"], } for item in distinct_values ], } ) if not distinct_values: continue chosen_value = ( distinct_values[0]["value"] if len(distinct_values) == 1 else normalize_merge_field_value(survivor_row[field_name]) ) field_should_lock = field_name in survivor_locks or (len(distinct_values) == 1 and bool(distinct_values[0]["locked_document_ids"])) if field_should_lock: fields_to_lock.add(field_name) if chosen_value is not None: builtin_field_updates[field_name] = chosen_value continue if ( field_name == "page_count" and len(distinct_values) == 1 and normalize_merge_field_value(survivor_row[field_name]) is None and chosen_value is not None ): builtin_field_updates[field_name] = chosen_value return { "document_ids": [int(row["id"]) for row in document_rows], "survivor_document_id": survivor_id, "loser_document_ids": loser_document_ids, "status": "blocked" if blockers else "ready", "blocking_conflicts": blockers, "machine_field_conflicts": machine_field_conflicts, "survivor_selection": survivor_selection_payload( connection, survivor_row, rule=resolved_selection_rule, ), "_custom_field_names": custom_fields, "_custom_field_updates": custom_field_updates, "_builtin_field_updates": builtin_field_updates, "_fields_to_lock": sorted(fields_to_lock), } def evaluate_reconcile_candidate_group( connection: sqlite3.Connection, document_rows: list[sqlite3.Row], ) -> dict[str, object]: custom_fields = reconcile_custom_field_names(connection) root_group = evaluate_document_merge_group( connection, document_rows, custom_fields=custom_fields, ) survivor_document_id = int(root_group["survivor_document_id"]) family_descriptors = { int(row["id"]): document_family_descriptor(connection, row) for row in document_rows } survivor_family = family_descriptors[survivor_document_id] child_merge_groups: list[dict[str, object]] = [] child_blockers: list[dict[str, object]] = [] for slot_index, survivor_child_row in enumerate(list(survivor_family["child_rows"])): child_group_rows = [ family_descriptors[int(row["id"])]["child_rows"][slot_index] for row in document_rows ] child_group = evaluate_document_merge_group( connection, child_group_rows, custom_fields=custom_fields, forced_survivor_document_id=int(survivor_child_row["id"]), selection_rule=f"family_slot[{slot_index}]", ) child_merge_groups.append(child_group) if child_group["status"] != "ready": child_blockers.append( { "type": "family_child_conflict", "slot_index": slot_index, "survivor_document_id": int(survivor_child_row["id"]), "child_document_ids": child_group["document_ids"], "blocking_conflicts": child_group["blocking_conflicts"], } ) blocking_conflicts = [*root_group["blocking_conflicts"], *child_blockers] root_group["content_hash"] = document_rows[0]["content_hash"] if document_rows else None root_group["family_fingerprint"] = survivor_family["fingerprint"] root_group["family_child_group_count"] = len(child_merge_groups) root_group["status"] = "blocked" if blocking_conflicts else "ready" root_group["blocking_conflicts"] = blocking_conflicts root_group["_child_merge_groups"] = child_merge_groups return root_group def resolve_reconcile_root_document_rows( connection: sqlite3.Connection, document_ids: list[int] | None, ) -> list[sqlite3.Row]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) if not normalized_document_ids: return [] root_rows_by_id: dict[int, sqlite3.Row] = {} ordered_root_ids: list[int] = [] for document_id in normalized_document_ids: root_row = get_document_family_root_row_for_assignment(connection, document_id) root_id = int(root_row["id"]) full_root_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (root_id,), ).fetchone() if full_root_row is None: raise RetrieverError(f"Unknown document id: {document_id}") if full_root_row["canonical_status"] != CANONICAL_STATUS_ACTIVE: raise RetrieverError( f"Document {document_id} cannot be reconciled because family root document {root_id} is not active." ) if root_id not in root_rows_by_id: root_rows_by_id[root_id] = full_root_row ordered_root_ids.append(root_id) return [root_rows_by_id[root_id] for root_id in ordered_root_ids] def find_reconcile_candidate_groups( connection: sqlite3.Connection, *, basis: str, candidate_root_document_ids: list[int] | None = None, ) -> list[dict[str, object]]: if basis != "content_hash": raise RetrieverError(f"Unsupported reconciliation basis: {basis}") normalized_candidate_root_ids = None if candidate_root_document_ids is not None: normalized_candidate_root_ids = list(dict.fromkeys(int(document_id) for document_id in candidate_root_document_ids)) if not normalized_candidate_root_ids: return [] candidate_id_clause = "" candidate_params: list[object] = [CANONICAL_STATUS_ACTIVE] if normalized_candidate_root_ids is not None: placeholders = ", ".join("?" for _ in normalized_candidate_root_ids) candidate_id_clause = f"\n AND id IN ({placeholders})" candidate_params.extend(normalized_candidate_root_ids) candidate_rows = connection.execute( f""" SELECT * FROM documents WHERE canonical_status = ? AND lifecycle_status = 'active' AND parent_document_id IS NULL AND content_hash IS NOT NULL AND text_status IN ('ok', 'partial') {candidate_id_clause} ORDER BY content_hash ASC, id ASC """, tuple(candidate_params), ).fetchall() groups_by_hash: dict[str, list[sqlite3.Row]] = defaultdict(list) for row in candidate_rows: groups_by_hash[str(row["content_hash"])].append(row) candidate_groups: list[dict[str, object]] = [] for content_hash, grouped_rows in sorted(groups_by_hash.items()): if len(grouped_rows) < 2: continue grouped_rows_by_family: dict[str, list[sqlite3.Row]] = defaultdict(list) for row in grouped_rows: family_descriptor = document_family_descriptor(connection, row) grouped_rows_by_family[str(family_descriptor["fingerprint"])].append(row) for family_fingerprint, family_rows in sorted(grouped_rows_by_family.items()): if len(family_rows) < 2: continue candidate_group = evaluate_reconcile_candidate_group(connection, family_rows) candidate_group["content_hash"] = content_hash candidate_group["family_fingerprint"] = family_fingerprint candidate_groups.append(candidate_group) return candidate_groups def rebind_document_dedupe_keys( connection: sqlite3.Connection, *, survivor_document_id: int, loser_document_id: int, ) -> None: rows = connection.execute( """ SELECT basis, key_value FROM document_dedupe_keys WHERE document_id = ? ORDER BY basis ASC, key_value ASC """, (loser_document_id,), ).fetchall() for row in rows: bind_document_dedupe_key( connection, basis=str(row["basis"]), key_value=str(row["key_value"]), document_id=survivor_document_id, ) connection.execute( "DELETE FROM document_dedupe_keys WHERE document_id = ?", (loser_document_id,), ) def transfer_manual_dataset_memberships( connection: sqlite3.Connection, *, survivor_document_id: int, loser_document_id: int, ) -> None: rows = connection.execute( """ SELECT dataset_id FROM dataset_documents WHERE document_id = ? AND dataset_source_id IS NULL ORDER BY dataset_id ASC """, (loser_document_id,), ).fetchall() for row in rows: ensure_dataset_document_membership( connection, dataset_id=int(row["dataset_id"]), document_id=survivor_document_id, dataset_source_id=None, ) def apply_evaluated_document_merge_group( connection: sqlite3.Connection, *, paths: dict[str, Path], merge_basis: str, merge_group: dict[str, object], ) -> dict[str, object]: survivor_document_id = int(merge_group["survivor_document_id"]) loser_document_ids = [int(document_id) for document_id in merge_group["loser_document_ids"]] custom_field_names = [str(field_name) for field_name in merge_group["_custom_field_names"]] custom_field_updates = dict(merge_group["_custom_field_updates"]) builtin_field_updates = dict(merge_group["_builtin_field_updates"]) fields_to_lock = list(merge_group["_fields_to_lock"]) pending_conflicts: list[tuple[int, sqlite3.Row, sqlite3.Row]] = [] merge_event_ids: list[int] = [] moved_occurrence_count = 0 for loser_document_id in loser_document_ids: survivor_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (survivor_document_id,), ).fetchone() loser_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (loser_document_id,), ).fetchone() if survivor_row is None or loser_row is None: raise RetrieverError(f"Unable to load merge pair survivor={survivor_document_id} loser={loser_document_id}.") merge_event_id = insert_document_merge_event( connection, survivor_document_id=survivor_document_id, loser_document_id=loser_document_id, merge_basis=merge_basis, pre_merge_survivor_snapshot=document_merge_snapshot(connection, survivor_document_id), pre_merge_loser_snapshot=document_merge_snapshot(connection, loser_document_id), artifact_counts={ "survivor": document_artifact_counts(connection, survivor_document_id), "loser": document_artifact_counts(connection, loser_document_id), }, ) merge_event_ids.append(merge_event_id) pending_conflicts.append((merge_event_id, survivor_row, loser_row)) transfer_manual_dataset_memberships( connection, survivor_document_id=survivor_document_id, loser_document_id=loser_document_id, ) rebind_document_dedupe_keys( connection, survivor_document_id=survivor_document_id, loser_document_id=loser_document_id, ) occurrence_rows = connection.execute( """ SELECT id FROM document_occurrences WHERE document_id = ? ORDER BY id ASC """, (loser_document_id,), ).fetchall() moved_occurrence_count += len(occurrence_rows) connection.execute( """ UPDATE document_occurrences SET document_id = ?, updated_at = ? WHERE document_id = ? """, (survivor_document_id, utc_now(), loser_document_id), ) connection.execute( "DELETE FROM dataset_documents WHERE document_id = ?", (loser_document_id,), ) connection.execute( "DELETE FROM document_control_number_aliases WHERE document_id IN (?, ?)", (survivor_document_id, loser_document_id), ) connection.execute( "DELETE FROM canonical_metadata_conflicts WHERE document_id = ?", (loser_document_id,), ) cleanup_document_artifacts(paths, connection, loser_row) delete_document_related_rows(connection, loser_document_id) connection.execute( """ UPDATE documents SET control_number = NULL, dataset_id = NULL, canonical_status = ?, merged_into_document_id = ?, rel_path = ?, lifecycle_status = 'deleted', updated_at = ? WHERE id = ? """, ( CANONICAL_STATUS_MERGED, survivor_document_id, tombstone_rel_path_for_merged_document(loser_document_id), utc_now(), loser_document_id, ), ) if custom_field_updates or builtin_field_updates or fields_to_lock: survivor_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (survivor_document_id,), ).fetchone() if survivor_row is None: raise RetrieverError(f"Unknown survivor document id: {survivor_document_id}") update_fields: dict[str, object] = {} update_fields.update(custom_field_updates) update_fields.update(builtin_field_updates) if fields_to_lock: update_fields[MANUAL_FIELD_LOCKS_COLUMN] = json.dumps(sorted(dict.fromkeys(fields_to_lock))) if update_fields: update_fields["updated_at"] = utc_now() assignments = ", ".join(f"{quote_identifier(field_name)} = ?" for field_name in update_fields) connection.execute( f"UPDATE documents SET {assignments} WHERE id = ?", [*update_fields.values(), survivor_document_id], ) refresh_source_backed_dataset_memberships_for_document(connection, survivor_document_id) refresh_document_from_occurrences(connection, survivor_document_id) final_survivor_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (survivor_document_id,), ).fetchone() if final_survivor_row is None: raise RetrieverError(f"Unknown survivor document id after merge: {survivor_document_id}") for merge_event_id, pre_merge_survivor_row, pre_merge_loser_row in pending_conflicts: insert_document_field_conflicts( connection, merge_event_id=merge_event_id, survivor_document_id=survivor_document_id, final_survivor_row=final_survivor_row, pre_merge_survivor_row=pre_merge_survivor_row, pre_merge_loser_row=pre_merge_loser_row, custom_field_names=custom_field_names, ) return { "survivor_document_id": survivor_document_id, "loser_document_ids": loser_document_ids, "merge_event_ids": merge_event_ids, "moved_occurrence_count": moved_occurrence_count, } def apply_reconcile_group( connection: sqlite3.Connection, *, paths: dict[str, Path], basis: str, candidate_group: dict[str, object], ) -> dict[str, object]: root_merge_result = apply_evaluated_document_merge_group( connection, paths=paths, merge_basis=basis, merge_group=candidate_group, ) child_merge_groups = list(candidate_group.get("_child_merge_groups") or []) child_merge_event_ids: list[int] = [] child_moved_occurrence_count = 0 for child_group in child_merge_groups: child_merge_result = apply_evaluated_document_merge_group( connection, paths=paths, merge_basis=f"{basis}:family_child", merge_group=child_group, ) child_merge_event_ids.extend(list(child_merge_result["merge_event_ids"])) child_moved_occurrence_count += int(child_merge_result["moved_occurrence_count"]) return { "content_hash": candidate_group["content_hash"], "document_ids": candidate_group["document_ids"], "survivor_document_id": root_merge_result["survivor_document_id"], "loser_document_ids": root_merge_result["loser_document_ids"], "status": "merged", "merge_event_ids": [*root_merge_result["merge_event_ids"], *child_merge_event_ids], "moved_occurrence_count": int(root_merge_result["moved_occurrence_count"]) + child_moved_occurrence_count, "blocking_conflicts": [], "machine_field_conflicts": candidate_group["machine_field_conflicts"], "survivor_selection": candidate_group["survivor_selection"], "family_child_group_count": len(child_merge_groups), } def reconcile_duplicates( root: Path, *, basis: str, apply_changes: bool, document_ids: list[int] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, dataset_names: list[str] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, ) -> dict[str, object]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) selector_inputs_present = bool(query.strip() or raw_bates or raw_filters or dataset_names or from_run_id is not None) if normalized_document_ids and (selector_inputs_present or select_from_scope): raise RetrieverError( "reconcile-duplicates accepts either --doc-id selectors or query/filter/scope selectors, not both." ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) selector_payload: dict[str, object] = {"mode": "all_documents"} selected_from_scope = False selected_document_ids: list[int] = [] selected_root_document_rows: list[sqlite3.Row] | None = None if normalized_document_ids: selected_document_ids = normalized_document_ids selected_root_document_rows = resolve_reconcile_root_document_rows(connection, normalized_document_ids) selector_payload = { "mode": "document_ids", "document_ids": selected_document_ids, "root_document_ids": [int(row["id"]) for row in selected_root_document_rows], } elif selector_inputs_present or select_from_scope: selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if not scope_run_selector_has_inputs(selector): raise RetrieverError( "No document selection active. Provide --doc-id or at least one of " "--keyword, --filter, --bates, --dataset, --from-run-id, or --select-from-scope." ) selected_document_ids, _, _ = resolve_seed_documents_for_scope_selector(connection, selector) selected_root_document_rows = resolve_reconcile_root_document_rows(connection, selected_document_ids) selector_payload = { "mode": "scope_search", "scope": selector, "document_ids": selected_document_ids, "root_document_ids": [int(row["id"]) for row in selected_root_document_rows], } selected_from_scope = bool(select_from_scope) selected_root_document_ids = ( None if selected_root_document_rows is None else [int(row["id"]) for row in selected_root_document_rows] ) selection_payload: dict[str, object] = { "selector": selector_payload, "selected_from_scope": selected_from_scope, } if selected_root_document_ids is not None: selection_payload["selected_document_ids"] = selected_document_ids selection_payload["selected_document_count"] = len(selected_document_ids) selection_payload["selected_root_document_ids"] = selected_root_document_ids selection_payload["selected_root_document_count"] = len(selected_root_document_ids) candidate_groups = find_reconcile_candidate_groups( connection, basis=basis, candidate_root_document_ids=selected_root_document_ids, ) if not apply_changes: return { "status": "ok", "basis": basis, "mode": "dry-run", **selection_payload, "candidate_group_count": len(candidate_groups), "mergeable_group_count": sum(1 for group in candidate_groups if group["status"] == "ready"), "blocked_group_count": sum(1 for group in candidate_groups if group["status"] == "blocked"), "candidate_groups": [ { key: value for key, value in group.items() if not key.startswith("_") } for group in candidate_groups ], } applied_groups: list[dict[str, object]] = [] blocked_groups: list[dict[str, object]] = [] connection.execute("BEGIN") try: for candidate_group in candidate_groups: if candidate_group["status"] != "ready": blocked_groups.append( { key: value for key, value in candidate_group.items() if not key.startswith("_") } ) continue applied_groups.append( apply_reconcile_group( connection, paths=paths, basis=basis, candidate_group=candidate_group, ) ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "basis": basis, "mode": "apply", **selection_payload, "candidate_group_count": len(candidate_groups), "merged_group_count": len(applied_groups), "blocked_group_count": len(blocked_groups), "applied_groups": applied_groups, "blocked_groups": blocked_groups, } finally: connection.close() def counted_delete( connection: sqlite3.Connection, *, count_sql: str, delete_sql: str, params: tuple[object, ...] = (), ) -> int: row = connection.execute(count_sql, params).fetchone() deleted_count = int(row[0] or 0) if row is not None else 0 connection.execute(delete_sql, params) return deleted_count def reset_auto_entity_graph(connection: sqlite3.Connection) -> dict[str, int]: auto_document_links_deleted = counted_delete( connection, count_sql="SELECT COUNT(*) FROM document_entities WHERE assignment_mode = 'auto'", delete_sql="DELETE FROM document_entities WHERE assignment_mode = 'auto'", ) auto_resolution_keys_deleted = counted_delete( connection, count_sql=""" SELECT COUNT(*) FROM entity_resolution_keys WHERE identifier_id IS NULL OR identifier_id IN ( SELECT id FROM entity_identifiers WHERE COALESCE(source_kind, 'auto') = 'auto' ) """, delete_sql=""" DELETE FROM entity_resolution_keys WHERE identifier_id IS NULL OR identifier_id IN ( SELECT id FROM entity_identifiers WHERE COALESCE(source_kind, 'auto') = 'auto' ) """, ) auto_identifiers_deleted = counted_delete( connection, count_sql="SELECT COUNT(*) FROM entity_identifiers WHERE COALESCE(source_kind, 'auto') = 'auto'", delete_sql="DELETE FROM entity_identifiers WHERE COALESCE(source_kind, 'auto') = 'auto'", ) auto_entities_deleted = counted_delete( connection, count_sql=""" SELECT COUNT(*) FROM entities WHERE entity_origin IN ('observed', 'identified') AND display_name_source = 'auto' AND canonical_status = 'active' AND NOT EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = entities.id AND COALESCE(ei.source_kind, 'auto') != 'auto' ) AND NOT EXISTS ( SELECT 1 FROM document_entities de WHERE de.entity_id = entities.id AND de.assignment_mode != 'auto' ) AND NOT EXISTS ( SELECT 1 FROM entity_resolution_keys erk WHERE erk.entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entity_overrides eo WHERE eo.source_entity_id = entities.id OR eo.replacement_entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entity_merge_blocks emb WHERE emb.left_entity_id = entities.id OR emb.right_entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entities merged_child WHERE merged_child.merged_into_entity_id = entities.id ) """, delete_sql=""" DELETE FROM entities WHERE entity_origin IN ('observed', 'identified') AND display_name_source = 'auto' AND canonical_status = 'active' AND NOT EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = entities.id AND COALESCE(ei.source_kind, 'auto') != 'auto' ) AND NOT EXISTS ( SELECT 1 FROM document_entities de WHERE de.entity_id = entities.id AND de.assignment_mode != 'auto' ) AND NOT EXISTS ( SELECT 1 FROM entity_resolution_keys erk WHERE erk.entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entity_overrides eo WHERE eo.source_entity_id = entities.id OR eo.replacement_entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entity_merge_blocks emb WHERE emb.left_entity_id = entities.id OR emb.right_entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entities merged_child WHERE merged_child.merged_into_entity_id = entities.id ) """, ) return { "auto_document_links_deleted": auto_document_links_deleted, "auto_resolution_keys_deleted": auto_resolution_keys_deleted, "auto_identifiers_deleted": auto_identifiers_deleted, "auto_entities_deleted": auto_entities_deleted, } def entity_rebuild_document_ids( connection: sqlite3.Connection, document_ids: list[int] | None, ) -> list[int]: if document_ids: normalized_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids)) placeholders = ",".join("?" for _ in normalized_ids) rows = connection.execute( f""" SELECT id FROM documents WHERE id IN ({placeholders}) ORDER BY id ASC """, tuple(normalized_ids), ).fetchall() found_ids = [int(row["id"]) for row in rows] found_id_set = set(found_ids) missing_ids = [document_id for document_id in normalized_ids if document_id not in found_id_set] if missing_ids: raise RetrieverError(f"Unknown document id(s): {', '.join(str(item) for item in missing_ids)}") return found_ids return [ int(row["id"]) for row in connection.execute( """ SELECT id FROM documents WHERE canonical_status != ? ORDER BY id ASC """, (CANONICAL_STATUS_MERGED,), ).fetchall() ] def entity_graph_counts(connection: sqlite3.Connection) -> dict[str, int]: return { "entity_count": int(connection.execute("SELECT COUNT(*) FROM entities").fetchone()[0] or 0), "active_entity_count": int( connection.execute( "SELECT COUNT(*) FROM entities WHERE canonical_status = ?", (ENTITY_STATUS_ACTIVE,), ).fetchone()[0] or 0 ), "document_entity_count": int(connection.execute("SELECT COUNT(*) FROM document_entities").fetchone()[0] or 0), "resolution_key_count": int(connection.execute("SELECT COUNT(*) FROM entity_resolution_keys").fetchone()[0] or 0), } ENTITY_REBUILD_ACTIVE_STATUSES = {"resetting", "planning", "rebuilding"} ENTITY_REBUILD_TERMINAL_STATUSES = {"completed", "canceled", "failed"} ENTITY_REBUILD_ITEM_STATUSES = ("pending", "leased", "committed", "failed", "cancelled") ENTITY_REBUILD_RESET_STAGES = ( "document_entities", "resolution_keys", "identifiers", "entities", "complete", ) ENTITY_REBUILD_LEASE_SECONDS = 45 def new_entity_rebuild_run_id(now: datetime | None = None) -> str: return new_ingest_v2_run_id(now) def entity_rebuild_run_row_by_id(connection: sqlite3.Connection, run_id: str) -> sqlite3.Row | None: return connection.execute( "SELECT * FROM entity_rebuild_runs WHERE run_id = ?", (run_id,), ).fetchone() def latest_entity_rebuild_run_row(connection: sqlite3.Connection) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM entity_rebuild_runs ORDER BY id DESC LIMIT 1 """ ).fetchone() def active_entity_rebuild_run_row(connection: sqlite3.Connection) -> sqlite3.Row | None: return connection.execute( f""" SELECT * FROM entity_rebuild_runs WHERE status IN ({", ".join("?" for _ in ENTITY_REBUILD_ACTIVE_STATUSES)}) AND cancel_requested_at IS NULL ORDER BY id ASC LIMIT 1 """, tuple(sorted(ENTITY_REBUILD_ACTIVE_STATUSES)), ).fetchone() def require_entity_rebuild_run_row(connection: sqlite3.Connection, run_id: str) -> sqlite3.Row: row = entity_rebuild_run_row_by_id(connection, run_id) if row is None: raise RetrieverError(f"Unknown entity rebuild run id: {run_id}") return row def entity_rebuild_conflict_payload(root: Path, active_row: sqlite3.Row, *, message: str) -> dict[str, object]: run_id = str(active_row["run_id"]) quoted_root = shlex.quote(str(root)) quoted_run_id = shlex.quote(run_id) return { "ok": False, "error": "active_entity_rebuild_run", "active_run_id": run_id, "message": message, "status_command": f"rebuild-entities-status {quoted_root} --run-id {quoted_run_id}", "cancel_command": f"rebuild-entities-cancel {quoted_root} --run-id {quoted_run_id}", } def raise_if_entity_rebuild_active(connection: sqlite3.Connection, root: Path, *, command_name: str) -> None: active_row = active_entity_rebuild_run_row(connection) if active_row is None: return raise RetrieverStructuredError( f"{command_name} cannot run while entity rebuild run {active_row['run_id']} is active.", entity_rebuild_conflict_payload( root, active_row, message=f"{command_name} cannot run while a resumable entity rebuild is active.", ), ) def entity_rebuild_status_counts(connection: sqlite3.Connection, *, run_id: str) -> dict[str, int]: counts = {status: 0 for status in ENTITY_REBUILD_ITEM_STATUSES} rows = connection.execute( """ SELECT status, COUNT(*) AS count FROM entity_rebuild_items WHERE run_id = ? GROUP BY status """, (run_id,), ).fetchall() for row in rows: counts[str(row["status"])] = int(row["count"] or 0) return counts def entity_rebuild_status_payload( connection: sqlite3.Connection, root: Path, row: sqlite3.Row, *, budget_seconds: int = DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, ) -> dict[str, object]: run_id = str(row["run_id"]) counts = entity_rebuild_status_counts(connection, run_id=run_id) reset_counts = decode_json_text(row["reset_counts_json"], default={}) or {} cursor = decode_json_text(row["cursor_json"], default={}) or {} totals_row = connection.execute( """ SELECT COALESCE(SUM(document_synced), 0) AS documents_synced, COALESCE(SUM(auto_links_created), 0) AS auto_links_created FROM entity_rebuild_items WHERE run_id = ? """, (run_id,), ).fetchone() root_arg = shlex.quote(str(root)) run_id_arg = shlex.quote(run_id) budget_arg = str(int(budget_seconds)) next_commands: list[str] if row["cancel_requested_at"] is not None or str(row["status"]) in ENTITY_REBUILD_TERMINAL_STATUSES: next_commands = [f"rebuild-entities-status {root_arg} --run-id {run_id_arg}"] else: next_commands = [ f"rebuild-entities-run-step {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}", ] return { "run_id": run_id, "mode": str(row["mode"]), "phase": str(row["phase"]), "status": str(row["status"]), "batch_size": int(row["batch_size"] or 0), "counts": {"work_items": counts}, "progress": { "reset_stage": str(row["reset_stage"] or ""), "reset_counts": reset_counts, "cursor": cursor, "documents_scanned": sum(counts.values()), "documents_synced": int(totals_row["documents_synced"] or 0) if totals_row is not None else 0, "auto_links_created": int(totals_row["auto_links_created"] or 0) if totals_row is not None else 0, }, "graph": entity_graph_counts(connection), "created_at": row["created_at"], "started_at": row["started_at"], "completed_at": row["completed_at"], "cancel_requested_at": row["cancel_requested_at"], "last_error": row["error"], "next_recommended_commands": next_commands, } def rebuild_entities_start( root: Path, *, document_ids: list[int] | None = None, batch_size: int = 500, budget_seconds: int | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) normalized_batch_size = max(1, min(int(batch_size or 500), 5000)) paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="rebuild-entities-start"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="rebuild-entities-start") active_row = active_entity_rebuild_run_row(connection) if active_row is not None: raise RetrieverStructuredError( f"Entity rebuild run {active_row['run_id']} is already active.", entity_rebuild_conflict_payload( root, active_row, message="A resumable entity rebuild is active in this workspace.", ), ) selected_document_ids = entity_rebuild_document_ids(connection, document_ids) if document_ids else [] full_rebuild = not document_ids run_id = new_entity_rebuild_run_id() now = utc_now() phase = "resetting" if full_rebuild else "planning" connection.execute("BEGIN") try: connection.execute( """ INSERT INTO entity_rebuild_runs ( run_id, mode, phase, status, document_ids_json, batch_size, reset_stage, reset_counts_json, cursor_json, created_at, started_at, last_heartbeat_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, "full" if full_rebuild else "selected", phase, phase, compact_json_text(selected_document_ids), normalized_batch_size, "document_entities", compact_json_text( { "auto_document_links_deleted": 0, "auto_resolution_keys_deleted": 0, "auto_identifiers_deleted": 0, "auto_entities_deleted": 0, } ), compact_json_text({"selected_offset": 0, "last_document_id": 0, "planned": 0}), now, now, now, ), ) connection.commit() except Exception: connection.rollback() raise row = require_entity_rebuild_run_row(connection, run_id) return {"ok": True, "created": True, **entity_rebuild_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() def rebuild_entities_status(root: Path, *, run_id: str | None = None, budget_seconds: int | None = None) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = require_entity_rebuild_run_row(connection, run_id) if run_id else latest_entity_rebuild_run_row(connection) if row is None: return { "ok": True, "status": "none", "phase": None, "run_id": None, "counts": {"work_items": {status: 0 for status in ENTITY_REBUILD_ITEM_STATUSES}}, "progress": { "reset_stage": None, "reset_counts": {}, "cursor": {}, "documents_scanned": 0, "documents_synced": 0, "auto_links_created": 0, }, "graph": entity_graph_counts(connection), "next_recommended_commands": [], } return {"ok": True, **entity_rebuild_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() def rebuild_entities_cancel(root: Path, *, run_id: str) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="rebuild-entities-cancel"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = require_entity_rebuild_run_row(connection, run_id) now = utc_now() connection.execute("BEGIN") try: if str(row["status"]) not in ENTITY_REBUILD_TERMINAL_STATUSES: connection.execute( """ UPDATE entity_rebuild_items SET status = 'cancelled', lease_owner = NULL, lease_expires_at = NULL, updated_at = ? WHERE run_id = ? AND status IN ('pending', 'leased') """, (now, run_id), ) connection.execute( """ UPDATE entity_rebuild_runs SET phase = 'canceled', status = 'canceled', cancel_requested_at = COALESCE(cancel_requested_at, ?), completed_at = COALESCE(completed_at, ?), last_heartbeat_at = ? WHERE run_id = ? """, (now, now, now, run_id), ) connection.commit() except Exception: connection.rollback() raise updated_row = require_entity_rebuild_run_row(connection, run_id) return {"ok": True, "cancelled": str(updated_row["status"]) == "canceled", **entity_rebuild_status_payload(connection, root, updated_row)} finally: connection.close() def entity_rebuild_reset_stage_delete_sql(stage: str) -> tuple[str, str]: if stage == "document_entities": return ( "auto_document_links_deleted", """ DELETE FROM document_entities WHERE id IN ( SELECT id FROM document_entities WHERE assignment_mode = 'auto' ORDER BY id ASC LIMIT ? ) """, ) if stage == "resolution_keys": return ( "auto_resolution_keys_deleted", """ DELETE FROM entity_resolution_keys WHERE id IN ( SELECT id FROM entity_resolution_keys WHERE identifier_id IS NULL OR identifier_id IN ( SELECT id FROM entity_identifiers WHERE COALESCE(source_kind, 'auto') = 'auto' ) ORDER BY id ASC LIMIT ? ) """, ) if stage == "identifiers": return ( "auto_identifiers_deleted", """ DELETE FROM entity_identifiers WHERE id IN ( SELECT id FROM entity_identifiers WHERE COALESCE(source_kind, 'auto') = 'auto' ORDER BY id ASC LIMIT ? ) """, ) if stage == "entities": return ( "auto_entities_deleted", """ DELETE FROM entities WHERE id IN ( SELECT id FROM entities WHERE entity_origin IN ('observed', 'identified') AND display_name_source = 'auto' AND canonical_status = 'active' AND NOT EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = entities.id AND COALESCE(ei.source_kind, 'auto') != 'auto' ) AND NOT EXISTS ( SELECT 1 FROM document_entities de WHERE de.entity_id = entities.id AND de.assignment_mode != 'auto' ) AND NOT EXISTS ( SELECT 1 FROM entity_resolution_keys erk WHERE erk.entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entity_overrides eo WHERE eo.source_entity_id = entities.id OR eo.replacement_entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entity_merge_blocks emb WHERE emb.left_entity_id = entities.id OR emb.right_entity_id = entities.id ) AND NOT EXISTS ( SELECT 1 FROM entities merged_child WHERE merged_child.merged_into_entity_id = entities.id ) ORDER BY id ASC LIMIT ? ) """, ) raise RetrieverError(f"Unsupported entity rebuild reset stage: {stage}") def entity_rebuild_next_reset_stage(stage: str) -> str: try: index = ENTITY_REBUILD_RESET_STAGES.index(stage) except ValueError: return "complete" return ENTITY_REBUILD_RESET_STAGES[min(index + 1, len(ENTITY_REBUILD_RESET_STAGES) - 1)] def entity_rebuild_reset_step(connection: sqlite3.Connection, *, run_id: str, batch_size: int) -> dict[str, object]: row = require_entity_rebuild_run_row(connection, run_id) stage = str(row["reset_stage"] or "document_entities") counts = decode_json_text(row["reset_counts_json"], default={}) or {} if stage == "complete": return {"stage": stage, "deleted": 0, "reset_complete": True} count_key, delete_sql = entity_rebuild_reset_stage_delete_sql(stage) now = utc_now() connection.execute("BEGIN") try: cursor = connection.execute(delete_sql, (max(1, int(batch_size)),)) deleted = int(cursor.rowcount or 0) if deleted: counts[count_key] = int(counts.get(count_key) or 0) + deleted next_stage = stage next_phase = "resetting" else: next_stage = entity_rebuild_next_reset_stage(stage) next_phase = "planning" if next_stage == "complete" else "resetting" connection.execute( """ UPDATE entity_rebuild_runs SET reset_stage = ?, reset_counts_json = ?, phase = ?, status = ?, last_heartbeat_at = ? WHERE run_id = ? AND status = 'resetting' AND cancel_requested_at IS NULL """, (next_stage, compact_json_text(counts), next_phase, next_phase, now, run_id), ) connection.commit() return { "stage": stage, "deleted": deleted, "reset_complete": next_stage == "complete", "next_stage": next_stage, } except Exception: connection.rollback() raise def entity_rebuild_plan_step(connection: sqlite3.Connection, *, run_id: str, batch_size: int) -> dict[str, object]: row = require_entity_rebuild_run_row(connection, run_id) cursor = decode_json_text(row["cursor_json"], default={}) or {} now = utc_now() planned_ids: list[int] = [] if str(row["mode"]) == "selected": selected_ids = [int(value) for value in list(decode_json_text(row["document_ids_json"], default=[]) or [])] offset = int(cursor.get("selected_offset") or 0) planned_ids = selected_ids[offset : offset + max(1, int(batch_size))] cursor["selected_offset"] = offset + len(planned_ids) planning_complete = int(cursor["selected_offset"]) >= len(selected_ids) else: last_document_id = int(cursor.get("last_document_id") or 0) rows = connection.execute( """ SELECT id FROM documents WHERE canonical_status != ? AND id > ? ORDER BY id ASC LIMIT ? """, (CANONICAL_STATUS_MERGED, last_document_id, max(1, int(batch_size))), ).fetchall() planned_ids = [int(item["id"]) for item in rows] if planned_ids: cursor["last_document_id"] = planned_ids[-1] planning_complete = len(planned_ids) < max(1, int(batch_size)) connection.execute("BEGIN") try: planned = 0 for ordinal, document_id in enumerate(planned_ids, start=int(cursor.get("planned") or 0) + 1): insert_cursor = connection.execute( """ INSERT OR IGNORE INTO entity_rebuild_items ( run_id, document_id, ordinal, status, created_at, updated_at ) VALUES (?, ?, ?, 'pending', ?, ?) """, (run_id, document_id, ordinal, now, now), ) planned += int(insert_cursor.rowcount or 0) cursor["planned"] = int(cursor.get("planned") or 0) + planned if planning_complete: item_count = int( connection.execute( "SELECT COUNT(*) FROM entity_rebuild_items WHERE run_id = ?", (run_id,), ).fetchone()[0] or 0 ) next_phase = "rebuilding" if item_count else "completed" completed_at = now if not item_count else None else: next_phase = "planning" completed_at = None connection.execute( """ UPDATE entity_rebuild_runs SET cursor_json = ?, phase = ?, status = ?, completed_at = COALESCE(completed_at, ?), last_heartbeat_at = ? WHERE run_id = ? AND phase = 'planning' AND cancel_requested_at IS NULL """, (compact_json_text(cursor), next_phase, next_phase, completed_at, now, run_id), ) connection.commit() return {"planned": planned, "planning_complete": planning_complete, "next_phase": next_phase} except Exception: connection.rollback() raise def entity_rebuild_reclaim_stale_items(connection: sqlite3.Connection, *, run_id: str) -> int: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) stale_cutoff = format_utc_timestamp(now_dt - timedelta(seconds=ENTITY_REBUILD_LEASE_SECONDS)) connection.execute("BEGIN") try: cursor = connection.execute( """ UPDATE entity_rebuild_items SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, updated_at = ? WHERE run_id = ? AND status = 'leased' AND lease_expires_at IS NOT NULL AND ( lease_expires_at <= ? OR updated_at <= ? ) """, (now, run_id, now, stale_cutoff), ) reclaimed = int(cursor.rowcount or 0) connection.commit() return reclaimed except Exception: connection.rollback() raise def entity_rebuild_claim_items( connection: sqlite3.Connection, *, run_id: str, worker_id: str, limit: int, ) -> list[sqlite3.Row]: now_dt = datetime.now(timezone.utc) now = format_utc_timestamp(now_dt) lease_expires_at = lease_expiration_after(ENTITY_REBUILD_LEASE_SECONDS, now=now_dt) connection.execute("BEGIN IMMEDIATE") try: row = require_entity_rebuild_run_row(connection, run_id) if str(row["phase"]) != "rebuilding" or row["cancel_requested_at"] is not None: connection.rollback() return [] claim_rows = connection.execute( """ SELECT id FROM entity_rebuild_items WHERE run_id = ? AND status = 'pending' ORDER BY ordinal ASC, id ASC LIMIT ? """, (run_id, max(1, int(limit))), ).fetchall() claim_ids = [int(item["id"]) for item in claim_rows] if claim_ids: placeholders = ",".join("?" for _ in claim_ids) connection.execute( f""" UPDATE entity_rebuild_items SET status = 'leased', lease_owner = ?, lease_expires_at = ?, attempts = attempts + 1, updated_at = ? WHERE run_id = ? AND status = 'pending' AND id IN ({placeholders}) """, (worker_id, lease_expires_at, now, run_id, *claim_ids), ) connection.execute( "UPDATE entity_rebuild_runs SET last_heartbeat_at = ? WHERE run_id = ?", (now, run_id), ) connection.commit() except Exception: connection.rollback() raise if not claim_ids: return [] placeholders = ",".join("?" for _ in claim_ids) return connection.execute( f""" SELECT * FROM entity_rebuild_items WHERE run_id = ? AND lease_owner = ? AND status = 'leased' AND id IN ({placeholders}) ORDER BY ordinal ASC, id ASC """, (run_id, worker_id, *claim_ids), ).fetchall() def entity_rebuild_release_items( connection: sqlite3.Connection, *, run_id: str, worker_id: str, item_ids: list[int], reason: str, ) -> int: if not item_ids: return 0 now = utc_now() placeholders = ",".join("?" for _ in item_ids) connection.execute("BEGIN") try: cursor = connection.execute( f""" UPDATE entity_rebuild_items SET status = 'pending', lease_owner = NULL, lease_expires_at = NULL, last_error = ?, updated_at = ? WHERE run_id = ? AND lease_owner = ? AND status = 'leased' AND id IN ({placeholders}) """, (reason, now, run_id, worker_id, *item_ids), ) released = int(cursor.rowcount or 0) connection.commit() return released except Exception: connection.rollback() raise def entity_rebuild_mark_item_failed( connection: sqlite3.Connection, *, run_id: str, item_id: int, worker_id: str, message: str, ) -> None: now = utc_now() connection.execute("BEGIN") try: connection.execute( """ UPDATE entity_rebuild_items SET status = 'failed', lease_owner = NULL, lease_expires_at = NULL, last_error = ?, updated_at = ? WHERE run_id = ? AND id = ? AND status = 'leased' AND lease_owner = ? """, (message, now, run_id, item_id, worker_id), ) connection.commit() except Exception: connection.rollback() raise def entity_rebuild_maybe_complete(connection: sqlite3.Connection, *, run_id: str) -> bool: row = require_entity_rebuild_run_row(connection, run_id) if str(row["phase"]) != "rebuilding" or row["cancel_requested_at"] is not None: return False remaining = int( connection.execute( """ SELECT COUNT(*) FROM entity_rebuild_items WHERE run_id = ? AND status IN ('pending', 'leased') """, (run_id,), ).fetchone()[0] or 0 ) if remaining: return False now = utc_now() connection.execute("BEGIN") try: connection.execute( """ UPDATE entity_rebuild_runs SET phase = 'completed', status = 'completed', completed_at = COALESCE(completed_at, ?), last_heartbeat_at = ? WHERE run_id = ? AND phase = 'rebuilding' AND cancel_requested_at IS NULL """, (now, now, run_id), ) connection.commit() return True except Exception: connection.rollback() raise def rebuild_entities_run_step( root: Path, *, run_id: str | None = None, budget_seconds: int | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) worker_id = ingest_v2_worker_id("entity-rebuild") executed_steps: list[str] = [] step_results: list[dict[str, object]] = [] with workspace_entity_rebuild_session(paths, command_name="rebuild-entities-run-step"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = require_entity_rebuild_run_row(connection, run_id) if run_id else active_entity_rebuild_run_row(connection) if row is None and run_id is None: row = latest_entity_rebuild_run_row(connection) if row is None: raise RetrieverError("No entity rebuild run exists in this workspace.") run_id = str(row["run_id"]) if str(row["status"]) in ENTITY_REBUILD_TERMINAL_STATUSES or row["cancel_requested_at"] is not None: return { "ok": True, "executed_steps": [], "reason": "run_terminal", "more_work_remaining": False, "run": entity_rebuild_status_payload(connection, root, row, budget_seconds=budget), } raise_if_ingest_v2_active(connection, root, command_name="rebuild-entities-run-step") while ingest_v2_deadline_remaining_seconds(deadline) >= INGEST_V2_RUN_STEP_MIN_REMAINING_SECONDS: row = require_entity_rebuild_run_row(connection, run_id) phase = str(row["phase"]) if phase == "resetting": result = entity_rebuild_reset_step(connection, run_id=run_id, batch_size=int(row["batch_size"] or 500)) executed_steps.append("reset") step_results.append(result) continue if phase == "planning": result = entity_rebuild_plan_step(connection, run_id=run_id, batch_size=int(row["batch_size"] or 500)) executed_steps.append("plan") step_results.append(result) continue if phase != "rebuilding": break reclaimed = entity_rebuild_reclaim_stale_items(connection, run_id=run_id) claim_limit = max(1, min(int(row["batch_size"] or 500), 100)) claimed_rows = entity_rebuild_claim_items( connection, run_id=run_id, worker_id=worker_id, limit=claim_limit, ) if not claimed_rows: completed = entity_rebuild_maybe_complete(connection, run_id=run_id) executed_steps.append("complete" if completed else "rebuild") step_results.append({"claimed": 0, "committed": 0, "failed": 0, "stale_reclaimed": reclaimed}) break committed = 0 failed = 0 processed_ids: set[int] = set() for item_row in claimed_rows: item_id = int(item_row["id"]) if ingest_v2_deadline_remaining_seconds(deadline) < INGEST_V2_RUN_STEP_MIN_REMAINING_SECONDS: break processed_ids.add(item_id) document_id = int(item_row["document_id"]) connection.execute("BEGIN") try: result = refresh_document_from_occurrences(connection, document_id) document_synced = 1 if result.get("canonical_status") == CANONICAL_STATUS_ACTIVE else 0 auto_links_created = 0 if document_synced: link_row = connection.execute( """ SELECT COUNT(*) FROM document_entities WHERE document_id = ? AND assignment_mode = 'auto' """, (document_id,), ).fetchone() auto_links_created = int(link_row[0] or 0) if link_row is not None else 0 now = utc_now() connection.execute( """ UPDATE entity_rebuild_items SET status = 'committed', lease_owner = NULL, lease_expires_at = NULL, document_synced = ?, auto_links_created = ?, last_error = NULL, updated_at = ? WHERE run_id = ? AND id = ? AND status = 'leased' AND lease_owner = ? """, (document_synced, auto_links_created, now, run_id, item_id, worker_id), ) connection.execute( "UPDATE entity_rebuild_runs SET last_heartbeat_at = ? WHERE run_id = ?", (now, run_id), ) connection.commit() committed += 1 except Exception as exc: rollback_open_transaction(connection) entity_rebuild_mark_item_failed( connection, run_id=run_id, item_id=item_id, worker_id=worker_id, message=f"{type(exc).__name__}: {exc}", ) failed += 1 unprocessed_ids = [int(item["id"]) for item in claimed_rows if int(item["id"]) not in processed_ids] released = entity_rebuild_release_items( connection, run_id=run_id, worker_id=worker_id, item_ids=unprocessed_ids, reason="Released because run-step budget was nearly exhausted.", ) entity_rebuild_maybe_complete(connection, run_id=run_id) executed_steps.append("rebuild") step_results.append( { "claimed": len(claimed_rows), "committed": committed, "failed": failed, "released": released, "stale_reclaimed": reclaimed, } ) updated_row = require_entity_rebuild_run_row(connection, run_id) reason = "run_terminal" if str(updated_row["status"]) in ENTITY_REBUILD_TERMINAL_STATUSES else "budget_exhausted" return { "ok": True, "executed_steps": executed_steps, "step_results": step_results, "reason": reason, "more_work_remaining": str(updated_row["status"]) not in ENTITY_REBUILD_TERMINAL_STATUSES, "run": entity_rebuild_status_payload(connection, root, updated_row, budget_seconds=budget), "remaining_budget_seconds": round(ingest_v2_deadline_remaining_seconds(deadline), 3), } finally: connection.close() def rebuild_entities( root: Path, *, document_ids: list[int] | None = None, batch_size: int = 500, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: normalized_batch_size = max(1, min(int(batch_size or 500), 5000)) paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="rebuild-entities") as rebuild_session: connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="rebuild-entities") raise_if_entity_rebuild_active(connection, root, command_name="rebuild-entities") if progress_callback is not None: progress_callback( { "status": "running", "phase": "preparing", "detail": "selecting documents", }, True, ) ids_to_rebuild = entity_rebuild_document_ids(connection, document_ids) full_rebuild = not document_ids total_documents = len(ids_to_rebuild) reset_counts = { "auto_document_links_deleted": 0, "auto_resolution_keys_deleted": 0, "auto_identifiers_deleted": 0, "auto_entities_deleted": 0, } if progress_callback is not None: progress_callback( { "status": "running" if total_documents or full_rebuild else "completed", "phase": "resetting" if full_rebuild else ("syncing" if total_documents else "completed"), "detail": f"0/{total_documents} documents", }, True, ) if full_rebuild: connection.execute("BEGIN") try: reset_counts = reset_auto_entity_graph(connection) connection.commit() except Exception: connection.rollback() raise documents_synced = 0 auto_links_created = 0 processed_documents = 0 if progress_callback is not None and full_rebuild and total_documents: progress_callback( { "status": "running", "phase": "syncing", "detail": f"0/{total_documents} documents", }, False, ) for offset in range(0, len(ids_to_rebuild), normalized_batch_size): batch_document_ids = ids_to_rebuild[offset : offset + normalized_batch_size] connection.execute("BEGIN") try: for document_id in batch_document_ids: result = refresh_document_from_occurrences(connection, document_id) if result.get("canonical_status") == CANONICAL_STATUS_ACTIVE: documents_synced += 1 row = connection.execute( """ SELECT COUNT(*) FROM document_entities WHERE document_id = ? AND assignment_mode = 'auto' """, (document_id,), ).fetchone() auto_links_created += int(row[0] or 0) if row is not None else 0 processed_documents += 1 if progress_callback is not None: progress_callback( { "status": "running", "phase": "syncing", "detail": f"{processed_documents}/{total_documents} documents", }, False, ) connection.commit() except Exception: connection.rollback() raise payload = { "status": "ok", "session_id": rebuild_session["id"], "mode": "full" if full_rebuild else "selected", "documents_scanned": total_documents, "documents_synced": documents_synced, "auto_links_created": auto_links_created, "batch_size": normalized_batch_size, **reset_counts, **entity_graph_counts(connection), } if progress_callback is not None: progress_callback( { "status": "completed", "phase": "completed", "detail": f"{processed_documents}/{total_documents} documents", }, True, ) return payload finally: connection.close() def extract_vault_filename_custodian_email_from_artifact(raw_value: object) -> str | None: text = normalize_entity_text(raw_value) marker_index = text.rfind("--") if marker_index < 0: return None suffix = text[marker_index + 2 :] dash_index = suffix.find("-") candidate = suffix[:dash_index] if dash_index >= 0 else suffix return normalize_entity_email(candidate) def active_entity_id_for_email_value( connection: sqlite3.Connection, email: str, *, exclude_entity_id: int | None = None, ) -> int | None: normalized_email = normalize_entity_email(email) if normalized_email is None: return None resolution_owner = active_entity_id_for_resolution_key( connection, { "identifier_type": "email", "normalized_value": normalized_email, }, ) if resolution_owner is not None and resolution_owner != exclude_entity_id: return resolution_owner row = connection.execute( """ SELECT e.id FROM entities e LEFT JOIN entity_identifiers ei ON ei.entity_id = e.id AND ei.identifier_type = 'email' AND ei.normalized_value = ? WHERE e.canonical_status = ? AND e.id != ? AND (e.primary_email = ? OR ei.id IS NOT NULL) ORDER BY CASE WHEN e.primary_email = ? THEN 0 ELSE 1 END, e.id ASC LIMIT 1 """, ( normalized_email, ENTITY_STATUS_ACTIVE, int(exclude_entity_id or 0), normalized_email, normalized_email, ), ).fetchone() return int(row["id"]) if row is not None else None def vault_filename_custodian_evidence_for_entity( connection: sqlite3.Connection, entity_id: int, ) -> dict[str, object]: link_rows = connection.execute( """ SELECT id, document_id, evidence_json FROM document_entities WHERE entity_id = ? AND role = 'custodian' ORDER BY id ASC """, (int(entity_id),), ).fetchall() document_ids: set[int] = set() occurrence_ids: set[int] = set() raw_values: set[str] = set() for row in link_rows: document_ids.add(int(row["document_id"])) try: evidence = json.loads(str(row["evidence_json"] or "{}")) except json.JSONDecodeError: evidence = {} if not isinstance(evidence, dict): evidence = {} raw_value = normalize_whitespace(str(evidence.get("raw_value") or "")) if raw_value: raw_values.add(raw_value) occurrence_id = evidence.get("occurrence_id") if occurrence_id is not None: try: occurrence_ids.add(int(occurrence_id)) except (TypeError, ValueError): pass if occurrence_ids: placeholders = ", ".join("?" for _ in occurrence_ids) occurrence_rows = connection.execute( f""" SELECT id, document_id, custodian FROM document_occurrences WHERE id IN ({placeholders}) ORDER BY id ASC """, tuple(sorted(occurrence_ids)), ).fetchall() for row in occurrence_rows: document_ids.add(int(row["document_id"])) raw_value = normalize_whitespace(str(row["custodian"] or "")) if raw_value: raw_values.add(raw_value) return { "document_ids": sorted(document_ids), "occurrence_ids": sorted(occurrence_ids), "raw_values": sorted(raw_values), "document_link_count": len(link_rows), } def cleaned_email_from_vault_filename_custodian_values(values: list[object]) -> str | None: for value in values: text = normalize_whitespace(str(value or "")) if not text: continue vault_parts = parse_google_vault_mbox_basename(Path(text).stem) if vault_parts is not None: return vault_parts["email"] for value in values: cleaned_email = extract_vault_filename_custodian_email_from_artifact(value) if cleaned_email: return cleaned_email return None def vault_filename_custodian_candidates(connection: sqlite3.Connection) -> list[dict[str, object]]: if not all(table_exists(connection, table_name) for table_name in ("entities", "entity_identifiers", "document_entities")): return [] rows = connection.execute( """ SELECT e.id, e.entity_type, e.display_name, e.primary_email, e.canonical_status FROM entities e WHERE e.canonical_status = ? AND ( COALESCE(e.primary_email, '') LIKE '%--%@%' OR EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = e.id AND ei.identifier_type = 'email' AND ei.normalized_value LIKE '%--%@%' ) ) ORDER BY e.id ASC """, (ENTITY_STATUS_ACTIVE,), ).fetchall() candidates: list[dict[str, object]] = [] for row in rows: entity_id = int(row["id"]) identifier_rows = connection.execute( """ SELECT id, identifier_type, display_value, normalized_value FROM entity_identifiers WHERE entity_id = ? ORDER BY id ASC """, (entity_id,), ).fetchall() polluted_email_values = [ normalize_whitespace(str(value or "")) for value in [row["primary_email"], *[item["normalized_value"] for item in identifier_rows if item["identifier_type"] == "email"]] if normalize_whitespace(str(value or "")) and "--" in normalize_whitespace(str(value or "")) and "@" in normalize_whitespace(str(value or "")) ] evidence = vault_filename_custodian_evidence_for_entity(connection, entity_id) raw_values = [str(value) for value in evidence["raw_values"]] cleaned_email = cleaned_email_from_vault_filename_custodian_values([*raw_values, *polluted_email_values]) if cleaned_email is None: continue target_entity_id = active_entity_id_for_email_value( connection, cleaned_email, exclude_entity_id=entity_id, ) action = "merge" if target_entity_id is not None else "rewrite" candidates.append( { "entity_id": entity_id, "entity_type": row["entity_type"], "display_name": row["display_name"], "polluted_email": polluted_email_values[0] if polluted_email_values else None, "cleaned_email": cleaned_email, "target_entity_id": target_entity_id, "action": action, "document_ids": evidence["document_ids"], "document_link_count": evidence["document_link_count"], "raw_values": raw_values, } ) return candidates def update_vault_filename_custodian_occurrences( connection: sqlite3.Connection, *, occurrence_ids: list[int], raw_values: list[str], cleaned_email: str, ) -> dict[str, object]: updated_occurrence_ids: set[int] = set() for occurrence_id in sorted({int(item) for item in occurrence_ids}): connection.execute( """ UPDATE document_occurrences SET custodian = ?, updated_at = ? WHERE id = ? AND COALESCE(custodian, '') != ? """, (cleaned_email, utc_now(), occurrence_id, cleaned_email), ) if int(connection.execute("SELECT changes()").fetchone()[0] or 0): updated_occurrence_ids.add(occurrence_id) normalized_raw_values = [normalize_whitespace(str(value or "")) for value in raw_values if normalize_whitespace(str(value or ""))] for raw_value in normalized_raw_values: matching_rows = connection.execute( """ SELECT id FROM document_occurrences WHERE custodian = ? ORDER BY id ASC """, (raw_value,), ).fetchall() matching_occurrence_ids = {int(row["id"]) for row in matching_rows} connection.execute( """ UPDATE document_occurrences SET custodian = ?, updated_at = ? WHERE custodian = ? """, (cleaned_email, utc_now(), raw_value), ) if int(connection.execute("SELECT changes()").fetchone()[0] or 0): updated_occurrence_ids.update(matching_occurrence_ids) affected_document_ids = [ int(row["document_id"]) for row in connection.execute( f""" SELECT DISTINCT document_id FROM document_occurrences WHERE id IN ({', '.join('?' for _ in updated_occurrence_ids)}) ORDER BY document_id ASC """, tuple(sorted(updated_occurrence_ids)), ).fetchall() ] if updated_occurrence_ids else [] return { "updated_occurrence_ids": sorted(updated_occurrence_ids), "affected_document_ids": affected_document_ids, } def delete_artifact_name_identifiers_for_entity(connection: sqlite3.Connection, entity_id: int) -> int: rows = connection.execute( """ SELECT id, display_value, normalized_value FROM entity_identifiers WHERE entity_id = ? AND identifier_type = 'name' ORDER BY id ASC """, (int(entity_id),), ).fetchall() deleted = 0 for row in rows: if not ( entity_name_identifier_looks_like_export_artifact(row["display_value"]) or entity_name_identifier_looks_like_export_artifact(row["normalized_value"]) ): continue connection.execute("DELETE FROM entity_resolution_keys WHERE identifier_id = ?", (int(row["id"]),)) connection.execute("DELETE FROM entity_identifiers WHERE id = ?", (int(row["id"]),)) deleted += int(connection.execute("SELECT changes()").fetchone()[0] or 0) return deleted def rewrite_vault_filename_custodian_entity( connection: sqlite3.Connection, *, entity_id: int, cleaned_email: str, ) -> dict[str, object]: deleted_artifact_names = delete_artifact_name_identifiers_for_entity(connection, entity_id) connection.execute( """ DELETE FROM entity_resolution_keys WHERE entity_id = ? AND key_type = 'email' AND normalized_value LIKE '%--%@%' """, (int(entity_id),), ) deleted_polluted_resolution_keys = int(connection.execute("SELECT changes()").fetchone()[0] or 0) connection.execute( """ UPDATE entity_identifiers SET display_value = ?, normalized_value = ?, is_verified = 1, updated_at = ? WHERE entity_id = ? AND identifier_type = 'email' AND normalized_value LIKE '%--%@%' """, (cleaned_email, cleaned_email, utc_now(), int(entity_id)), ) updated_email_identifiers = int(connection.execute("SELECT changes()").fetchone()[0] or 0) if updated_email_identifiers == 0: ensure_entity_identifier( connection, entity_id=int(entity_id), identifier={ "identifier_type": "email", "display_value": cleaned_email, "normalized_value": cleaned_email, "is_verified": 1, }, ) updated_email_identifiers = 1 email_identifier_rows = connection.execute( """ SELECT * FROM entity_identifiers WHERE entity_id = ? AND identifier_type = 'email' AND normalized_value = ? ORDER BY id ASC """, (int(entity_id), cleaned_email), ).fetchall() created_resolution_keys = 0 for identifier_row in email_identifier_rows: resolution_key_id = ensure_entity_resolution_key( connection, entity_id=int(entity_id), identifier_id=int(identifier_row["id"]), identifier={ "identifier_type": "email", "normalized_value": cleaned_email, }, ) if resolution_key_id is not None: created_resolution_keys += 1 connection.execute( """ UPDATE entities SET entity_type = ?, display_name = NULL, display_name_source = ?, updated_at = ? WHERE id = ? """, (ENTITY_TYPE_UNKNOWN, ENTITY_DISPLAY_SOURCE_AUTO, utc_now(), int(entity_id)), ) recompute_entity_caches(connection, int(entity_id)) return { "updated_email_identifiers": updated_email_identifiers, "deleted_artifact_name_identifiers": deleted_artifact_names, "deleted_polluted_resolution_keys": deleted_polluted_resolution_keys, "created_resolution_keys": created_resolution_keys, } def merge_vault_filename_custodian_entity( connection: sqlite3.Connection, *, loser_entity_id: int, survivor_entity_id: int, ) -> dict[str, object]: if int(loser_entity_id) == int(survivor_entity_id): raise RetrieverError("Cannot merge an entity into itself.") affected_document_ids = [ int(row["document_id"]) for row in connection.execute( """ SELECT DISTINCT document_id FROM document_entities WHERE entity_id IN (?, ?) ORDER BY document_id ASC """, (int(loser_entity_id), int(survivor_entity_id)), ).fetchall() ] duplicate_link_count = counted_delete( connection, count_sql=""" SELECT COUNT(*) FROM document_entities WHERE entity_id = ? AND EXISTS ( SELECT 1 FROM document_entities survivor_link WHERE survivor_link.document_id = document_entities.document_id AND survivor_link.role = document_entities.role AND survivor_link.entity_id = ? ) """, delete_sql=""" DELETE FROM document_entities WHERE entity_id = ? AND EXISTS ( SELECT 1 FROM document_entities survivor_link WHERE survivor_link.document_id = document_entities.document_id AND survivor_link.role = document_entities.role AND survivor_link.entity_id = ? ) """, params=(int(loser_entity_id), int(survivor_entity_id)), ) connection.execute( """ UPDATE document_entities SET entity_id = ?, updated_at = ? WHERE entity_id = ? """, (int(survivor_entity_id), utc_now(), int(loser_entity_id)), ) moved_link_count = int(connection.execute("SELECT changes()").fetchone()[0] or 0) connection.execute( """ UPDATE entity_overrides SET replacement_entity_id = ?, updated_at = ? WHERE replacement_entity_id = ? """, (int(survivor_entity_id), utc_now(), int(loser_entity_id)), ) connection.execute( """ UPDATE entity_overrides SET source_entity_id = ?, updated_at = ? WHERE source_entity_id = ? """, (int(survivor_entity_id), utc_now(), int(loser_entity_id)), ) deleted_artifact_names = delete_artifact_name_identifiers_for_entity(connection, loser_entity_id) connection.execute("DELETE FROM entity_resolution_keys WHERE entity_id = ?", (int(loser_entity_id),)) deleted_resolution_keys = int(connection.execute("SELECT changes()").fetchone()[0] or 0) connection.execute( """ UPDATE entities SET canonical_status = ?, merged_into_entity_id = ?, updated_at = ? WHERE id = ? """, (ENTITY_STATUS_MERGED, int(survivor_entity_id), utc_now(), int(loser_entity_id)), ) created_resolution_keys = ensure_manual_email_resolution_keys(connection, int(survivor_entity_id)) recompute_entity_caches(connection, int(survivor_entity_id)) refresh_documents_after_entity_graph_change(connection, affected_document_ids) return { "affected_document_ids": affected_document_ids, "moved_document_links": moved_link_count, "deduped_document_links": duplicate_link_count, "deleted_artifact_name_identifiers": deleted_artifact_names, "deleted_polluted_resolution_keys": deleted_resolution_keys, "created_resolution_keys": created_resolution_keys, } def purge_vault_filename_custodians( root: Path, *, apply: bool = False, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) command_name = "purge-vault-filename-custodians" with workspace_entity_rebuild_session(paths, command_name=command_name): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name=command_name) raise_if_entity_rebuild_active(connection, root, command_name=command_name) candidates = vault_filename_custodian_candidates(connection) if not apply: return { "status": "ok", "dry_run": True, "candidate_count": len(candidates), "candidates": candidates, } applied: list[dict[str, object]] = [] connection.execute("BEGIN") try: for candidate in candidates: entity_id = int(candidate["entity_id"]) cleaned_email = str(candidate["cleaned_email"]) evidence = vault_filename_custodian_evidence_for_entity(connection, entity_id) occurrence_result = update_vault_filename_custodian_occurrences( connection, occurrence_ids=[int(item) for item in evidence["occurrence_ids"]], raw_values=[str(item) for item in evidence["raw_values"]], cleaned_email=cleaned_email, ) target_entity_id = candidate.get("target_entity_id") if target_entity_id is not None: action_result = merge_vault_filename_custodian_entity( connection, loser_entity_id=entity_id, survivor_entity_id=int(target_entity_id), ) action = "merged" else: action_result = rewrite_vault_filename_custodian_entity( connection, entity_id=entity_id, cleaned_email=cleaned_email, ) refresh_documents_after_entity_graph_change( connection, [ *[int(item) for item in evidence["document_ids"]], *[int(item) for item in occurrence_result["affected_document_ids"]], ], ) action = "rewritten" applied.append( { **candidate, "action": action, **occurrence_result, **action_result, } ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "dry_run": False, "candidate_count": len(candidates), "applied_count": len(applied), "candidates": applied, **entity_graph_counts(connection), } finally: connection.close() def serialize_entity_identifier(row: sqlite3.Row) -> dict[str, object]: payload: dict[str, object] = { "id": int(row["id"]), "identifier_type": row["identifier_type"], "display_value": row["display_value"], "normalized_value": row["normalized_value"], "is_primary": bool(row["is_primary"]), "is_verified": bool(row["is_verified"]), "source_kind": row["source_kind"], } for key in ( "provider", "provider_scope", "identifier_name", "identifier_scope", "normalized_full_name", "normalized_sort_name", ): if payload_has_meaningful_value(row[key]): payload[key] = row[key] return payload def entity_identifiers_by_entity_id( connection: sqlite3.Connection, entity_ids: list[int], ) -> dict[int, list[dict[str, object]]]: if not entity_ids: return {} placeholders = ",".join("?" for _ in entity_ids) rows = connection.execute( f""" SELECT * FROM entity_identifiers WHERE entity_id IN ({placeholders}) ORDER BY entity_id ASC, is_primary DESC, is_verified DESC, id ASC """, tuple(entity_ids), ).fetchall() grouped: dict[int, list[dict[str, object]]] = defaultdict(list) for row in rows: grouped[int(row["entity_id"])].append(serialize_entity_identifier(row)) return grouped def serialize_entity_summary( row: sqlite3.Row, identifiers: list[dict[str, object]], ) -> dict[str, object]: roles = sorted({role for role in str(row["roles"] or "").split(",") if role}) emails = [ str(identifier["normalized_value"]) for identifier in identifiers if identifier.get("identifier_type") == "email" ] phones = [ str(identifier["display_value"]) for identifier in identifiers if identifier.get("identifier_type") == "phone" ] names = [ str(identifier["display_value"]) for identifier in identifiers if identifier.get("identifier_type") == "name" ] return { "id": int(row["id"]), "label": entity_display_label_from_row(row), "entity_type": row["entity_type"], "entity_origin": row["entity_origin"], "canonical_status": row["canonical_status"], "display_name": row["display_name"], "primary_email": row["primary_email"], "primary_phone": row["primary_phone"], "sort_name": row["sort_name"], "document_count": int(row["document_count"] or 0), "roles": roles, "emails": emails, "phones": phones, "names": names, } ENTITY_LIST_SORT_EXPRESSIONS = { "id": "e.id", "label": "LOWER(COALESCE(e.sort_name, e.display_name, e.primary_email, e.primary_phone, ''))", "display_name": "LOWER(COALESCE(e.display_name, ''))", "primary_email": "LOWER(COALESCE(e.primary_email, ''))", "primary_phone": "LOWER(COALESCE(e.primary_phone, ''))", "sort_name": "LOWER(COALESCE(e.sort_name, ''))", "entity_type": "e.entity_type", "entity_origin": "e.entity_origin", "canonical_status": "e.canonical_status", "entity_status": "e.canonical_status", "document_count": "document_count", } def normalize_entity_list_sort_field(raw_field: str | None) -> str: field_name = normalize_inline_whitespace(str(raw_field or "")).lower() if field_name == "status": field_name = "entity_status" if field_name == "type": field_name = "entity_type" if field_name == "origin": field_name = "entity_origin" if field_name == "email": field_name = "primary_email" if field_name not in ENTITY_LIST_SORT_EXPRESSIONS: allowed = ", ".join(sorted(ENTITY_LIST_SORT_EXPRESSIONS)) raise RetrieverError(f"Unsupported entity sort field: {raw_field}. Sortable fields: {allowed}.") return field_name def normalize_entity_list_sort_specs( *, sort: str | None = None, order: str | None = None, sort_specs: list[tuple[str, str]] | None = None, ) -> list[tuple[str, str]]: if sort_specs is not None: normalized_specs: list[tuple[str, str]] = [] for raw_field, raw_direction in sort_specs: field_name = normalize_entity_list_sort_field(raw_field) direction = normalize_inline_whitespace(str(raw_direction or "asc")).lower() if direction not in {"asc", "desc"}: raise RetrieverError("Sort direction must be 'asc' or 'desc'.") normalized_specs.append((field_name, direction)) return normalized_specs or [("document_count", "desc"), ("label", "asc"), ("id", "asc")] if sort: direction = normalize_inline_whitespace(str(order or "asc")).lower() if direction not in {"asc", "desc"}: raise RetrieverError("Sort direction must be 'asc' or 'desc'.") return [(normalize_entity_list_sort_field(sort), direction)] return [("document_count", "desc"), ("label", "asc"), ("id", "asc")] def entity_list_sort_spec_text(sort_specs: list[tuple[str, str]]) -> str: return ", ".join(f"{field_name} {direction}" for field_name, direction in sort_specs) def entity_list_order_sql(sort_specs: list[tuple[str, str]]) -> str: effective_specs = list(sort_specs) if not any(field_name == "id" for field_name, _ in effective_specs): effective_specs.append(("id", "asc")) parts: list[str] = [] for field_name, direction in effective_specs: expression = ENTITY_LIST_SORT_EXPRESSIONS[normalize_entity_list_sort_field(field_name)] parts.append(f"{expression} {direction.upper()}") return ", ".join(parts) def list_entities( root: Path, *, query: str | None = None, limit: int = 50, offset: int = 0, sort: str | None = None, order: str | None = None, sort_specs: list[tuple[str, str]] | None = None, include_ignored: bool = False, ) -> dict[str, object]: normalized_limit = max(1, min(int(limit or 50), 200)) normalized_offset = max(0, int(offset or 0)) normalized_query = normalize_whitespace(str(query or "")).lower() normalized_sort_specs = normalize_entity_list_sort_specs(sort=sort, order=order, sort_specs=sort_specs) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) where_clauses = ["e.canonical_status != ?"] if include_ignored else ["e.canonical_status = ?"] params: list[object] = [ENTITY_STATUS_MERGED if include_ignored else ENTITY_STATUS_ACTIVE] if normalized_query: like_query = f"%{normalized_query}%" where_clauses.append( """ ( LOWER(COALESCE(e.display_name, '')) LIKE ? OR LOWER(COALESCE(e.primary_email, '')) LIKE ? OR LOWER(COALESCE(e.primary_phone, '')) LIKE ? OR EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = e.id AND ( LOWER(COALESCE(ei.display_value, '')) LIKE ? OR LOWER(COALESCE(ei.normalized_value, '')) LIKE ? ) ) ) """ ) params.extend([like_query, like_query, like_query, like_query, like_query]) where_sql = " AND ".join(where_clauses) total_row = connection.execute( f""" SELECT COUNT(*) AS total FROM entities e WHERE {where_sql} """, tuple(params), ).fetchone() total_entities = int(total_row["total"] if total_row is not None else 0) rows = connection.execute( f""" SELECT e.*, COUNT(DISTINCT de.document_id) AS document_count, GROUP_CONCAT(DISTINCT de.role) AS roles FROM entities e LEFT JOIN document_entities de ON de.entity_id = e.id WHERE {where_sql} GROUP BY e.id ORDER BY {entity_list_order_sql(normalized_sort_specs)} LIMIT ? OFFSET ? """, (*params, normalized_limit, normalized_offset), ).fetchall() entity_ids = [int(row["id"]) for row in rows] identifiers_by_entity = entity_identifiers_by_entity_id(connection, entity_ids) entities = [ serialize_entity_summary(row, identifiers_by_entity.get(int(row["id"]), [])) for row in rows ] return { "status": "ok", "query": normalized_query, "limit": normalized_limit, "offset": normalized_offset, "total_hits": total_entities, "total": total_entities, "has_more": normalized_offset + len(entities) < total_entities, "next_offset": normalized_offset + normalized_limit if normalized_offset + len(entities) < total_entities else None, "sort": normalized_sort_specs[0][0], "order": normalized_sort_specs[0][1], "sort_spec": entity_list_sort_spec_text(normalized_sort_specs), "sort_override": serialize_sort_specs(normalized_sort_specs), "include_ignored": include_ignored, "entities": entities, **entity_graph_counts(connection), } finally: connection.close() def show_entity(root: Path, entity_id: int, *, document_limit: int = 25) -> dict[str, object]: normalized_document_limit = max(1, min(int(document_limit or 25), 200)) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = connection.execute( """ SELECT e.*, COUNT(DISTINCT de.document_id) AS document_count, GROUP_CONCAT(DISTINCT de.role) AS roles FROM entities e LEFT JOIN document_entities de ON de.entity_id = e.id WHERE e.id = ? GROUP BY e.id """, (int(entity_id),), ).fetchone() if row is None: raise RetrieverError(f"Unknown entity id: {entity_id}") identifiers = entity_identifiers_by_entity_id(connection, [int(entity_id)]).get(int(entity_id), []) role_counts = [ {"role": role_row["role"], "document_count": int(role_row["document_count"] or 0)} for role_row in connection.execute( """ SELECT role, COUNT(DISTINCT document_id) AS document_count FROM document_entities WHERE entity_id = ? GROUP BY role ORDER BY role ASC """, (int(entity_id),), ).fetchall() ] document_rows = connection.execute( """ SELECT de.role, de.ordinal, de.assignment_mode, de.observed_title, d.id AS document_id, d.control_number, d.rel_path, d.title, d.date_created FROM document_entities de JOIN documents d ON d.id = de.document_id WHERE de.entity_id = ? ORDER BY de.role ASC, de.ordinal ASC, d.id ASC LIMIT ? """, (int(entity_id), normalized_document_limit), ).fetchall() documents = [ { "document_id": int(document_row["document_id"]), "role": document_row["role"], "ordinal": int(document_row["ordinal"] or 0), "assignment_mode": document_row["assignment_mode"], "control_number": document_row["control_number"], "rel_path": document_row["rel_path"], "title": document_row["title"], "date_created": document_row["date_created"], "observed_title": document_row["observed_title"], } for document_row in document_rows ] return { "status": "ok", "entity": serialize_entity_summary(row, identifiers), "identifiers": identifiers, "role_counts": role_counts, "documents": documents, "document_limit": normalized_document_limit, } finally: connection.close() def entity_payload_by_id(connection: sqlite3.Connection, entity_id: int) -> dict[str, object]: row = connection.execute( """ SELECT e.*, COUNT(DISTINCT de.document_id) AS document_count, GROUP_CONCAT(DISTINCT de.role) AS roles FROM entities e LEFT JOIN document_entities de ON de.entity_id = e.id WHERE e.id = ? GROUP BY e.id """, (int(entity_id),), ).fetchone() if row is None: raise RetrieverError(f"Unknown entity id: {entity_id}") identifiers = entity_identifiers_by_entity_id(connection, [int(entity_id)]).get(int(entity_id), []) return { "entity": serialize_entity_summary(row, identifiers), "identifiers": identifiers, } def active_entity_row(connection: sqlite3.Connection, entity_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM entities WHERE id = ? """, (int(entity_id),), ).fetchone() if row is None: raise RetrieverError(f"Unknown entity id: {entity_id}") if row["canonical_status"] != ENTITY_STATUS_ACTIVE: raise RetrieverError(f"Entity {entity_id} is not active; status is {row['canonical_status']}.") return row def normalize_entity_type_arg(raw_entity_type: object) -> str: normalized = normalize_entity_lookup_text(raw_entity_type).replace("-", "_").replace(" ", "_") aliases = { "shared": ENTITY_TYPE_SHARED_MAILBOX, "shared_mailbox": ENTITY_TYPE_SHARED_MAILBOX, "system": ENTITY_TYPE_SYSTEM_MAILBOX, "system_mailbox": ENTITY_TYPE_SYSTEM_MAILBOX, "mailbox": ENTITY_TYPE_SHARED_MAILBOX, } entity_type = aliases.get(normalized, normalized) if entity_type not in ENTITY_TYPES: raise RetrieverError( f"Unsupported entity type: {raw_entity_type}. Supported types: {', '.join(sorted(ENTITY_TYPES))}." ) return entity_type def parse_manual_handle_identifier_arg(raw_value: object) -> dict[str, object]: parts = [part.strip() for part in str(raw_value or "").split(":", 2)] if len(parts) != 3: raise RetrieverError("Handle identifiers must use provider:scope:handle, e.g. slack:workspace:@jane.") provider = normalize_entity_identifier_name(parts[0]) provider_scope = normalize_entity_lookup_text(parts[1]) handle = normalize_entity_handle(parts[2]) if not provider or not provider_scope or not handle: raise RetrieverError("Handle identifiers require non-empty provider, scope, and handle values.") return { "identifier_type": "handle", "display_value": parts[2], "normalized_value": handle, "provider": provider, "provider_scope": provider_scope, "source_kind": "manual", } def parse_manual_external_id_identifier_arg(raw_value: object) -> dict[str, object]: parts = [part.strip() for part in str(raw_value or "").split(":", 2)] if len(parts) == 2: raw_name, raw_value_part = parts raw_scope = None elif len(parts) == 3: raw_name, raw_scope, raw_value_part = parts else: raise RetrieverError("External identifiers must use name:value or name:scope:value.") identifier_name = normalize_entity_identifier_name(raw_name) identifier_scope = normalize_entity_lookup_text(raw_scope) if raw_scope else None normalized_value = normalize_entity_lookup_text(raw_value_part) if not identifier_name or not normalized_value: raise RetrieverError("External identifiers require non-empty name and value fields.") return { "identifier_type": "external_id", "display_value": raw_value_part, "normalized_value": normalized_value, "identifier_name": identifier_name, "identifier_scope": identifier_scope, "source_kind": "manual", } def manual_entity_identifier_payloads( *, emails: list[str] | None = None, phones: list[str] | None = None, names: list[str] | None = None, handles: list[str] | None = None, external_ids: list[str] | None = None, ) -> list[dict[str, object]]: payloads: list[dict[str, object]] = [] for index, raw_email in enumerate(emails or []): email = normalize_entity_email(raw_email) if not email: raise RetrieverError(f"Invalid entity email identifier: {raw_email}") payloads.append( { "identifier_type": "email", "display_value": email, "normalized_value": email, "is_primary": 1 if index == 0 else 0, "is_verified": 1, "source_kind": "manual", } ) for index, raw_phone in enumerate(phones or []): phone = normalize_entity_phone(raw_phone) if phone is None: raise RetrieverError(f"Invalid entity phone identifier: {raw_phone}") payloads.append( { "identifier_type": "phone", "display_value": phone["display_value"], "normalized_value": phone["normalized_value"], "parsed_phone_json": json.dumps(phone["parsed_phone"], ensure_ascii=True, sort_keys=True), "is_primary": 1 if index == 0 else 0, "source_kind": "manual", } ) for index, raw_name in enumerate(names or []): parsed_name = parse_entity_name(raw_name) if parsed_name is None: raise RetrieverError(f"Invalid entity name identifier: {raw_name}") payloads.append( { "identifier_type": "name", "display_value": parsed_name["display_value"], "normalized_value": parsed_name["normalized_value"], "parsed_name_json": json.dumps(parsed_name["parsed_name"], ensure_ascii=True, sort_keys=True), "normalized_full_name": parsed_name["normalized_full_name"], "normalized_sort_name": parsed_name["normalized_sort_name"], "is_primary": 1 if index == 0 or parsed_name["is_full_name"] else 0, "source_kind": "manual", } ) payloads.extend(parse_manual_handle_identifier_arg(raw_handle) for raw_handle in handles or []) payloads.extend(parse_manual_external_id_identifier_arg(raw_external_id) for raw_external_id in external_ids or []) deduped_payloads: list[dict[str, object]] = [] seen_keys: set[str] = set() for payload in payloads: key = entity_candidate_identifier_key(payload) if key in seen_keys: continue seen_keys.add(key) deduped_payloads.append(payload) return deduped_payloads def manual_identifier_error_label(identifier: dict[str, object]) -> str: identifier_type = str(identifier.get("identifier_type") or "identifier") display_value = str(identifier.get("display_value") or identifier.get("normalized_value") or "") if identifier_type == "external_id" and identifier.get("identifier_name"): return f"{identifier.get('identifier_name')}:{display_value}" if identifier_type == "handle" and identifier.get("provider") and identifier.get("provider_scope"): return f"{identifier.get('provider')}:{identifier.get('provider_scope')}:{display_value}" return f"{identifier_type}:{display_value}" def assert_manual_resolution_identifiers_available( connection: sqlite3.Connection, identifiers: list[dict[str, object]], *, entity_id: int | None = None, ) -> None: for identifier in identifiers: if identifier.get("identifier_type") not in {"email", "handle", "external_id"}: continue owner_entity_id = active_entity_id_for_resolution_key(connection, identifier) if owner_entity_id is not None and (entity_id is None or owner_entity_id != int(entity_id)): raise RetrieverError( f"Identifier {manual_identifier_error_label(identifier)!r} already resolves to entity {owner_entity_id}." ) def ensure_manual_entity_identifiers( connection: sqlite3.Connection, *, entity_id: int, identifiers: list[dict[str, object]], ) -> dict[str, object]: identifier_ids: list[int] = [] created_count = 0 existing_count = 0 for identifier in identifiers: before_count = int(connection.execute("SELECT COUNT(*) FROM entity_identifiers").fetchone()[0] or 0) identifier_id = ensure_entity_identifier(connection, entity_id=int(entity_id), identifier=identifier) after_count = int(connection.execute("SELECT COUNT(*) FROM entity_identifiers").fetchone()[0] or 0) identifier_ids.append(identifier_id) if after_count > before_count: created_count += 1 else: existing_count += 1 return { "identifier_ids": identifier_ids, "created_identifier_count": created_count, "existing_identifier_count": existing_count, } def create_entity( root: Path, *, entity_type: str = ENTITY_TYPE_PERSON, display_name: str | None = None, notes: str | None = None, emails: list[str] | None = None, phones: list[str] | None = None, names: list[str] | None = None, handles: list[str] | None = None, external_ids: list[str] | None = None, ) -> dict[str, object]: normalized_entity_type = normalize_entity_type_arg(entity_type) normalized_display_name = normalize_whitespace(str(display_name or "")) or None normalized_notes = normalize_whitespace(str(notes or "")) or None identifiers = manual_entity_identifier_payloads( emails=emails, phones=phones, names=names, handles=handles, external_ids=external_ids, ) if normalized_display_name is None and not identifiers: raise RetrieverError("create-entity requires --display-name or at least one identifier.") paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="create-entity"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="create-entity") raise_if_entity_rebuild_active(connection, root, command_name="create-entity") connection.execute("BEGIN") try: assert_manual_resolution_identifiers_available(connection, identifiers) connection.execute( """ INSERT INTO entities ( entity_type, display_name, notes, display_name_source, entity_origin, canonical_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( normalized_entity_type, normalized_display_name, normalized_notes, ENTITY_DISPLAY_SOURCE_MANUAL if normalized_display_name else ENTITY_DISPLAY_SOURCE_AUTO, ENTITY_ORIGIN_MANUAL, ENTITY_STATUS_ACTIVE, utc_now(), utc_now(), ), ) entity_id = int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) identifier_counts = ensure_manual_entity_identifiers( connection, entity_id=entity_id, identifiers=identifiers, ) created_resolution_keys = ensure_manual_email_resolution_keys(connection, entity_id) recompute_entity_caches(connection, entity_id) payload = entity_payload_by_id(connection, entity_id) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "created": True, "entity_id": entity_id, "created_resolution_keys": created_resolution_keys, **identifier_counts, **payload, } finally: connection.close() def edit_entity( root: Path, entity_id: int, *, entity_type: str | None = None, display_name: str | None = None, clear_display_name: bool = False, notes: str | None = None, clear_notes: bool = False, add_emails: list[str] | None = None, add_phones: list[str] | None = None, add_names: list[str] | None = None, add_handles: list[str] | None = None, add_external_ids: list[str] | None = None, ) -> dict[str, object]: if clear_display_name and display_name is not None: raise RetrieverError("Use either --display-name or --clear-display-name, not both.") if clear_notes and notes is not None: raise RetrieverError("Use either --notes or --clear-notes, not both.") normalized_entity_type = normalize_entity_type_arg(entity_type) if entity_type is not None else None normalized_display_name = normalize_whitespace(str(display_name or "")) if display_name is not None else None normalized_notes = normalize_whitespace(str(notes or "")) if notes is not None else None identifiers = manual_entity_identifier_payloads( emails=add_emails, phones=add_phones, names=add_names, handles=add_handles, external_ids=add_external_ids, ) has_entity_update = ( normalized_entity_type is not None or display_name is not None or clear_display_name or notes is not None or clear_notes ) if not has_entity_update and not identifiers: raise RetrieverError("edit-entity requires an entity field change or at least one added identifier.") paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="edit-entity"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="edit-entity") raise_if_entity_rebuild_active(connection, root, command_name="edit-entity") active_entity_row(connection, entity_id) affected_document_ids = [ int(row["document_id"]) for row in connection.execute( """ SELECT DISTINCT document_id FROM document_entities WHERE entity_id = ? ORDER BY document_id ASC """, (int(entity_id),), ).fetchall() ] connection.execute("BEGIN") try: assert_manual_resolution_identifiers_available(connection, identifiers, entity_id=int(entity_id)) update_values: dict[str, object] = { "entity_origin": ENTITY_ORIGIN_MANUAL, "updated_at": utc_now(), } if normalized_entity_type is not None: update_values["entity_type"] = normalized_entity_type if display_name is not None: update_values["display_name"] = normalized_display_name or None update_values["display_name_source"] = ENTITY_DISPLAY_SOURCE_MANUAL elif clear_display_name: update_values["display_name"] = None update_values["display_name_source"] = ENTITY_DISPLAY_SOURCE_AUTO if notes is not None: update_values["notes"] = normalized_notes or None elif clear_notes: update_values["notes"] = None set_clause = ", ".join(f"{quote_identifier(column)} = ?" for column in update_values) connection.execute( f""" UPDATE entities SET {set_clause} WHERE id = ? """, [*update_values.values(), int(entity_id)], ) identifier_counts = ensure_manual_entity_identifiers( connection, entity_id=int(entity_id), identifiers=identifiers, ) created_resolution_keys = ensure_manual_email_resolution_keys(connection, int(entity_id)) recompute_entity_caches(connection, int(entity_id)) refresh_documents_after_entity_graph_change(connection, affected_document_ids) payload = entity_payload_by_id(connection, int(entity_id)) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "entity_id": int(entity_id), "created_resolution_keys": created_resolution_keys, "affected_document_ids": affected_document_ids, **identifier_counts, **payload, } finally: connection.close() def entity_pair_key(left_entity_id: int, right_entity_id: int) -> tuple[int, int]: left = int(left_entity_id) right = int(right_entity_id) if left == right: raise RetrieverError("Entity pair requires two distinct entity ids.") return (left, right) if left < right else (right, left) def entity_merge_block_exists(connection: sqlite3.Connection, left_entity_id: int, right_entity_id: int) -> bool: left, right = entity_pair_key(left_entity_id, right_entity_id) row = connection.execute( """ SELECT 1 FROM entity_merge_blocks WHERE left_entity_id = ? AND right_entity_id = ? LIMIT 1 """, (left, right), ).fetchone() return row is not None def block_entity_merge( root: Path, left_entity_id: int, right_entity_id: int, *, reason: str | None = None, ) -> dict[str, object]: left, right = entity_pair_key(left_entity_id, right_entity_id) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="block-entity-merge") active_entity_row(connection, left) active_entity_row(connection, right) normalized_reason = normalize_whitespace(str(reason or "")) or None connection.execute("BEGIN") try: connection.execute( """ INSERT OR IGNORE INTO entity_merge_blocks ( left_entity_id, right_entity_id, reason, created_at ) VALUES (?, ?, ?, ?) """, (left, right, normalized_reason, utc_now()), ) inserted = connection.execute("SELECT changes()").fetchone()[0] connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "left_entity_id": left, "right_entity_id": right, "created": bool(inserted), "reason": normalized_reason, } finally: connection.close() def entity_profile_from_summary(entity: dict[str, object], identifiers: list[dict[str, object]]) -> dict[str, object]: emails = { normalize_entity_email(identifier.get("normalized_value") or identifier.get("display_value")) for identifier in identifiers if identifier.get("identifier_type") == "email" } emails = {email for email in emails if email} email_locals = {email.split("@", 1)[0] for email in emails} email_domains = {email.split("@", 1)[1] for email in emails} full_names = { normalize_entity_lookup_text(identifier.get("normalized_full_name") or identifier.get("display_value")) for identifier in identifiers if identifier.get("identifier_type") == "name" } full_names = {name for name in full_names if name and len(name.split()) >= 2} sort_names = { normalize_entity_lookup_text(identifier.get("normalized_sort_name")) for identifier in identifiers if identifier.get("identifier_type") == "name" } sort_names = {name for name in sort_names if name} phones = { normalize_entity_lookup_text(identifier.get("normalized_value") or identifier.get("display_value")) for identifier in identifiers if identifier.get("identifier_type") == "phone" } phones = {phone for phone in phones if phone} handles = { ( normalize_entity_lookup_text(identifier.get("provider")), normalize_entity_lookup_text(identifier.get("provider_scope")), normalize_entity_lookup_text(identifier.get("normalized_value")), ) for identifier in identifiers if identifier.get("identifier_type") == "handle" } handles = {handle for handle in handles if handle[2]} label = normalize_entity_lookup_text(entity.get("label")) display_name = normalize_entity_lookup_text(entity.get("display_name")) if display_name and len(display_name.split()) >= 2: full_names.add(display_name) if label and len(label.split()) >= 2 and "@" not in label: full_names.add(label) return { "emails": emails, "email_locals": email_locals, "email_domains": email_domains, "full_names": full_names, "sort_names": sort_names, "phones": phones, "handles": handles, "label": label, } def name_initial_family_pairs(full_names: set[str]) -> set[tuple[str, str]]: pairs: set[tuple[str, str]] = set() for name in full_names: parts = name.split() if len(parts) < 2: continue pairs.add((parts[-1], parts[0][:1])) return pairs def entity_similarity_reasons(left_profile: dict[str, object], right_profile: dict[str, object]) -> list[dict[str, object]]: reasons: list[dict[str, object]] = [] left_emails = set(left_profile["emails"]) # type: ignore[arg-type] right_emails = set(right_profile["emails"]) # type: ignore[arg-type] if left_emails & right_emails: reasons.append({"kind": "exact_email", "score": 100, "value": sorted(left_emails & right_emails)[0]}) left_full_names = set(left_profile["full_names"]) # type: ignore[arg-type] right_full_names = set(right_profile["full_names"]) # type: ignore[arg-type] if left_full_names & right_full_names: reasons.append({"kind": "exact_full_name", "score": 80, "value": sorted(left_full_names & right_full_names)[0]}) left_sort_names = set(left_profile["sort_names"]) # type: ignore[arg-type] right_sort_names = set(right_profile["sort_names"]) # type: ignore[arg-type] if left_sort_names & right_sort_names: reasons.append({"kind": "exact_sort_name", "score": 80, "value": sorted(left_sort_names & right_sort_names)[0]}) left_phones = set(left_profile["phones"]) # type: ignore[arg-type] right_phones = set(right_profile["phones"]) # type: ignore[arg-type] if left_phones & right_phones: reasons.append({"kind": "same_phone", "score": 70, "value": sorted(left_phones & right_phones)[0]}) left_handles = set(left_profile["handles"]) # type: ignore[arg-type] right_handles = set(right_profile["handles"]) # type: ignore[arg-type] shared_handles = left_handles & right_handles if shared_handles: handle = sorted(shared_handles)[0] reasons.append({"kind": "same_handle", "score": 65, "value": handle[2]}) left_initial_pairs = name_initial_family_pairs(left_full_names) right_initial_pairs = name_initial_family_pairs(right_full_names) if left_initial_pairs & right_initial_pairs: family, initial = sorted(left_initial_pairs & right_initial_pairs)[0] reasons.append({"kind": "family_name_and_initial", "score": 45, "value": f"{family}, {initial}"}) left_email_locals = set(left_profile["email_locals"]) # type: ignore[arg-type] right_email_locals = set(right_profile["email_locals"]) # type: ignore[arg-type] if left_email_locals & right_email_locals: reasons.append({"kind": "same_email_local_part", "score": 40, "value": sorted(left_email_locals & right_email_locals)[0]}) left_email_domains = set(left_profile["email_domains"]) # type: ignore[arg-type] right_email_domains = set(right_profile["email_domains"]) # type: ignore[arg-type] if left_email_domains & right_email_domains and (left_full_names & right_full_names or left_initial_pairs & right_initial_pairs): reasons.append({"kind": "same_email_domain", "score": 20, "value": sorted(left_email_domains & right_email_domains)[0]}) return reasons def load_active_entity_summaries( connection: sqlite3.Connection, *, query: str | None = None, limit: int = 500, ) -> tuple[list[dict[str, object]], dict[int, list[dict[str, object]]]]: normalized_limit = max(1, min(int(limit or 500), 5000)) normalized_query = normalize_whitespace(str(query or "")).lower() where_clauses = ["e.canonical_status = ?"] params: list[object] = [ENTITY_STATUS_ACTIVE] if normalized_query: like_query = f"%{normalized_query}%" where_clauses.append( """ ( LOWER(COALESCE(e.display_name, '')) LIKE ? OR LOWER(COALESCE(e.primary_email, '')) LIKE ? OR LOWER(COALESCE(e.primary_phone, '')) LIKE ? OR EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = e.id AND ( LOWER(COALESCE(ei.display_value, '')) LIKE ? OR LOWER(COALESCE(ei.normalized_value, '')) LIKE ? ) ) ) """ ) params.extend([like_query, like_query, like_query, like_query, like_query]) rows = connection.execute( f""" SELECT e.*, COUNT(DISTINCT de.document_id) AS document_count, GROUP_CONCAT(DISTINCT de.role) AS roles FROM entities e LEFT JOIN document_entities de ON de.entity_id = e.id WHERE {' AND '.join(where_clauses)} GROUP BY e.id ORDER BY document_count DESC, COALESCE(e.sort_name, e.display_name, e.primary_email, e.primary_phone, '') ASC, e.id ASC LIMIT ? """, (*params, normalized_limit), ).fetchall() entity_ids = [int(row["id"]) for row in rows] identifiers_by_entity = entity_identifiers_by_entity_id(connection, entity_ids) return ( [serialize_entity_summary(row, identifiers_by_entity.get(int(row["id"]), [])) for row in rows], identifiers_by_entity, ) def similar_entities(root: Path, entity_id: int, *, limit: int = 25) -> dict[str, object]: normalized_limit = max(1, min(int(limit or 25), 200)) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) active_entity_row(connection, entity_id) entities, identifiers_by_entity = load_active_entity_summaries(connection, limit=5000) summaries_by_id = {int(entity["id"]): entity for entity in entities} target = summaries_by_id.get(int(entity_id)) if target is None: raise RetrieverError(f"Entity {entity_id} is not active.") profiles = { int(entity["id"]): entity_profile_from_summary( entity, identifiers_by_entity.get(int(entity["id"]), []), ) for entity in entities } target_profile = profiles[int(entity_id)] suggestions: list[dict[str, object]] = [] for candidate in entities: candidate_id = int(candidate["id"]) if candidate_id == int(entity_id): continue if entity_merge_block_exists(connection, int(entity_id), candidate_id): continue reasons = entity_similarity_reasons(target_profile, profiles[candidate_id]) if not reasons: continue score = sum(int(reason["score"]) for reason in reasons) suggestions.append( { "entity": candidate, "score": score, "reasons": reasons, } ) suggestions.sort( key=lambda item: ( -int(item["score"]), -int(item["entity"]["document_count"]), # type: ignore[index] str(item["entity"]["label"]), # type: ignore[index] int(item["entity"]["id"]), # type: ignore[index] ) ) return { "status": "ok", "entity": target, "suggestions": suggestions[:normalized_limit], "limit": normalized_limit, } finally: connection.close() def resolution_key_matches_row(connection: sqlite3.Connection, row: sqlite3.Row, *, exclude_id: int | None = None) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM entity_resolution_keys WHERE key_type = ? AND COALESCE(provider, '') = COALESCE(?, '') AND COALESCE(provider_scope, '') = COALESCE(?, '') AND COALESCE(identifier_name, '') = COALESCE(?, '') AND COALESCE(identifier_scope, '') = COALESCE(?, '') AND normalized_value = ? AND (? IS NULL OR id != ?) ORDER BY id ASC LIMIT 1 """, ( row["key_type"], row["provider"], row["provider_scope"], row["identifier_name"], row["identifier_scope"], row["normalized_value"], exclude_id, exclude_id, ), ).fetchone() def matching_survivor_identifier_id( connection: sqlite3.Connection, *, survivor_entity_id: int, identifier_row: sqlite3.Row, ) -> int | None: row = connection.execute( """ SELECT id FROM entity_identifiers WHERE entity_id = ? AND identifier_type = ? AND normalized_value = ? AND COALESCE(provider, '') = COALESCE(?, '') AND COALESCE(provider_scope, '') = COALESCE(?, '') AND COALESCE(identifier_name, '') = COALESCE(?, '') AND COALESCE(identifier_scope, '') = COALESCE(?, '') ORDER BY id ASC LIMIT 1 """, ( survivor_entity_id, identifier_row["identifier_type"], identifier_row["normalized_value"], identifier_row["provider"], identifier_row["provider_scope"], identifier_row["identifier_name"], identifier_row["identifier_scope"], ), ).fetchone() return int(row["id"]) if row is not None else None def identifier_payload_from_row(row: sqlite3.Row) -> dict[str, object]: return { "identifier_type": row["identifier_type"], "display_value": row["display_value"], "normalized_value": row["normalized_value"], "provider": row["provider"], "provider_scope": row["provider_scope"], "identifier_name": row["identifier_name"], "identifier_scope": row["identifier_scope"], "parsed_name_json": row["parsed_name_json"], "parsed_phone_json": row["parsed_phone_json"], "normalized_full_name": row["normalized_full_name"], "normalized_sort_name": row["normalized_sort_name"], "is_primary": row["is_primary"], "is_verified": row["is_verified"], "source_kind": row["source_kind"], } def move_loser_identifiers_to_survivor( connection: sqlite3.Connection, *, loser_entity_id: int, survivor_entity_id: int, ) -> dict[str, int]: moved_identifiers = 0 deduped_identifiers = 0 moved_resolution_keys = 0 deleted_resolution_keys = 0 identifier_rows = connection.execute( """ SELECT * FROM entity_identifiers WHERE entity_id = ? ORDER BY id ASC """, (loser_entity_id,), ).fetchall() for identifier_row in identifier_rows: loser_identifier_id = int(identifier_row["id"]) target_identifier_id = matching_survivor_identifier_id( connection, survivor_entity_id=survivor_entity_id, identifier_row=identifier_row, ) if target_identifier_id is None: connection.execute( """ UPDATE entity_identifiers SET entity_id = ?, source_kind = 'manual', updated_at = ? WHERE id = ? """, (survivor_entity_id, utc_now(), loser_identifier_id), ) target_identifier_id = loser_identifier_id moved_identifiers += 1 else: connection.execute( """ UPDATE entity_identifiers SET source_kind = 'manual', updated_at = ? WHERE id = ? """, (utc_now(), target_identifier_id), ) deduped_identifiers += 1 key_rows = connection.execute( """ SELECT * FROM entity_resolution_keys WHERE identifier_id = ? ORDER BY id ASC """, (loser_identifier_id,), ).fetchall() for key_row in key_rows: existing_key = resolution_key_matches_row(connection, key_row, exclude_id=int(key_row["id"])) if existing_key is not None: connection.execute("DELETE FROM entity_resolution_keys WHERE id = ?", (int(key_row["id"]),)) deleted_resolution_keys += 1 else: connection.execute( """ UPDATE entity_resolution_keys SET entity_id = ?, identifier_id = ?, updated_at = ? WHERE id = ? """, (survivor_entity_id, target_identifier_id, utc_now(), int(key_row["id"])), ) moved_resolution_keys += 1 if target_identifier_id != loser_identifier_id: connection.execute("DELETE FROM entity_identifiers WHERE id = ?", (loser_identifier_id,)) for key_row in connection.execute( """ SELECT * FROM entity_resolution_keys WHERE entity_id = ? ORDER BY id ASC """, (loser_entity_id,), ).fetchall(): existing_key = resolution_key_matches_row(connection, key_row, exclude_id=int(key_row["id"])) if existing_key is not None: connection.execute("DELETE FROM entity_resolution_keys WHERE id = ?", (int(key_row["id"]),)) deleted_resolution_keys += 1 else: connection.execute( """ UPDATE entity_resolution_keys SET entity_id = ?, updated_at = ? WHERE id = ? """, (survivor_entity_id, utc_now(), int(key_row["id"])), ) moved_resolution_keys += 1 connection.execute( """ UPDATE entity_identifiers SET source_kind = 'manual', updated_at = ? WHERE entity_id = ? """, (utc_now(), survivor_entity_id), ) return { "moved_identifiers": moved_identifiers, "deduped_identifiers": deduped_identifiers, "moved_resolution_keys": moved_resolution_keys, "deleted_resolution_keys": deleted_resolution_keys, } def ensure_manual_email_resolution_keys(connection: sqlite3.Connection, entity_id: int) -> int: created = 0 for identifier_row in connection.execute( """ SELECT * FROM entity_identifiers WHERE entity_id = ? AND identifier_type IN ('email', 'handle', 'external_id') ORDER BY id ASC """, (entity_id,), ).fetchall(): before = connection.execute("SELECT COUNT(*) FROM entity_resolution_keys").fetchone()[0] ensure_entity_resolution_key( connection, entity_id=entity_id, identifier_id=int(identifier_row["id"]), identifier=identifier_payload_from_row(identifier_row), ) after = connection.execute("SELECT COUNT(*) FROM entity_resolution_keys").fetchone()[0] created += max(0, int(after or 0) - int(before or 0)) return created def refresh_documents_after_entity_graph_change(connection: sqlite3.Connection, document_ids: list[int]) -> None: for document_id in sorted(set(int(item) for item in document_ids)): rebuild_document_entity_caches(connection, document_id) refresh_documents_fts_row(connection, document_id) def merge_entities( root: Path, source_entity_id: int, target_entity_id: int, *, force: bool = False, reason: str | None = None, ) -> dict[str, object]: loser_entity_id = int(source_entity_id) survivor_entity_id = int(target_entity_id) if loser_entity_id == survivor_entity_id: raise RetrieverError("Cannot merge an entity into itself.") paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="merge-entities"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="merge-entities") raise_if_entity_rebuild_active(connection, root, command_name="merge-entities") loser_row = active_entity_row(connection, loser_entity_id) survivor_row = active_entity_row(connection, survivor_entity_id) if entity_merge_block_exists(connection, loser_entity_id, survivor_entity_id) and not force: raise RetrieverError("These entities have a merge block. Pass --force to merge anyway.") affected_document_ids = [ int(row["document_id"]) for row in connection.execute( """ SELECT DISTINCT document_id FROM document_entities WHERE entity_id IN (?, ?) ORDER BY document_id ASC """, (loser_entity_id, survivor_entity_id), ).fetchall() ] connection.execute("BEGIN") try: identifier_counts = move_loser_identifiers_to_survivor( connection, loser_entity_id=loser_entity_id, survivor_entity_id=survivor_entity_id, ) duplicate_link_count = counted_delete( connection, count_sql=""" SELECT COUNT(*) FROM document_entities WHERE entity_id = ? AND EXISTS ( SELECT 1 FROM document_entities survivor_link WHERE survivor_link.document_id = document_entities.document_id AND survivor_link.role = document_entities.role AND survivor_link.entity_id = ? ) """, delete_sql=""" DELETE FROM document_entities WHERE entity_id = ? AND EXISTS ( SELECT 1 FROM document_entities survivor_link WHERE survivor_link.document_id = document_entities.document_id AND survivor_link.role = document_entities.role AND survivor_link.entity_id = ? ) """, params=(loser_entity_id, survivor_entity_id), ) connection.execute( """ UPDATE document_entities SET entity_id = ?, updated_at = ? WHERE entity_id = ? """, (survivor_entity_id, utc_now(), loser_entity_id), ) moved_link_count = int(connection.execute("SELECT changes()").fetchone()[0] or 0) connection.execute( """ UPDATE entity_overrides SET replacement_entity_id = ?, updated_at = ? WHERE replacement_entity_id = ? """, (survivor_entity_id, utc_now(), loser_entity_id), ) connection.execute( """ UPDATE entity_overrides SET source_entity_id = ?, updated_at = ? WHERE source_entity_id = ? """, (survivor_entity_id, utc_now(), loser_entity_id), ) connection.execute( """ UPDATE entities SET entity_origin = ?, updated_at = ? WHERE id = ? """, (ENTITY_ORIGIN_MANUAL, utc_now(), survivor_entity_id), ) connection.execute( """ UPDATE entities SET canonical_status = ?, merged_into_entity_id = ?, updated_at = ? WHERE id = ? """, (ENTITY_STATUS_MERGED, survivor_entity_id, utc_now(), loser_entity_id), ) created_resolution_keys = ensure_manual_email_resolution_keys(connection, survivor_entity_id) recompute_entity_caches(connection, survivor_entity_id) refresh_documents_after_entity_graph_change(connection, affected_document_ids) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "source_entity_id": loser_entity_id, "target_entity_id": survivor_entity_id, "source_label": entity_display_label_from_row(loser_row), "target_label": entity_display_label_from_row(survivor_row), "force": bool(force), "reason": normalize_whitespace(str(reason or "")) or None, "affected_document_ids": affected_document_ids, "moved_document_links": moved_link_count, "deduped_document_links": duplicate_link_count, "created_resolution_keys": created_resolution_keys, **identifier_counts, } finally: connection.close() def ignore_override_keys_for_entity(connection: sqlite3.Connection, entity_id: int) -> list[dict[str, object]]: keys: dict[tuple[str | None, str | None, str | None], dict[str, object]] = {} rows = connection.execute( """ SELECT role, evidence_json FROM document_entities WHERE entity_id = ? ORDER BY id ASC """, (int(entity_id),), ).fetchall() for row in rows: try: evidence = json.loads(str(row["evidence_json"] or "{}")) except json.JSONDecodeError: evidence = {} if not isinstance(evidence, dict): evidence = {} role = str(row["role"] or "") raw_value = normalize_whitespace(str(evidence.get("raw_value") or "")) candidate_key = normalize_whitespace(str(evidence.get("normalized_candidate_key") or "")) if candidate_key: keys[(role, candidate_key, raw_value or None)] = { "role": role, "normalized_candidate_key": candidate_key, "source_hint": raw_value or None, } if raw_value: for candidate in parse_entity_candidates(raw_value, role=role): parsed_key = normalize_whitespace(str(candidate.get("normalized_candidate_key") or "")) if parsed_key: keys[(role, parsed_key, raw_value)] = { "role": role, "normalized_candidate_key": parsed_key, "source_hint": raw_value, } return list(keys.values()) def ignore_entity(root: Path, entity_id: int, *, reason: str | None = None) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="ignore-entity"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="ignore-entity") raise_if_entity_rebuild_active(connection, root, command_name="ignore-entity") entity_row = active_entity_row(connection, entity_id) affected_document_ids = [ int(row["document_id"]) for row in connection.execute( """ SELECT DISTINCT document_id FROM document_entities WHERE entity_id = ? ORDER BY document_id ASC """, (int(entity_id),), ).fetchall() ] override_keys = ignore_override_keys_for_entity(connection, int(entity_id)) normalized_reason = normalize_whitespace(str(reason or "")) or None connection.execute("BEGIN") try: for override in override_keys: connection.execute( """ INSERT OR IGNORE INTO entity_overrides ( scope_type, scope_id, role, source_entity_id, normalized_candidate_key, replacement_entity_id, override_effect, source_hint, reason, created_at, updated_at ) VALUES ('global', NULL, ?, ?, ?, NULL, 'ignore', ?, ?, ?, ?) """, ( override.get("role"), int(entity_id), override.get("normalized_candidate_key"), override.get("source_hint"), normalized_reason, utc_now(), utc_now(), ), ) resolution_keys_deleted = counted_delete( connection, count_sql="SELECT COUNT(*) FROM entity_resolution_keys WHERE entity_id = ?", delete_sql="DELETE FROM entity_resolution_keys WHERE entity_id = ?", params=(int(entity_id),), ) document_links_deleted = counted_delete( connection, count_sql="SELECT COUNT(*) FROM document_entities WHERE entity_id = ?", delete_sql="DELETE FROM document_entities WHERE entity_id = ?", params=(int(entity_id),), ) connection.execute( """ UPDATE entities SET canonical_status = ?, merged_into_entity_id = NULL, updated_at = ? WHERE id = ? """, (ENTITY_STATUS_IGNORED, utc_now(), int(entity_id)), ) refresh_documents_after_entity_graph_change(connection, affected_document_ids) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "entity_id": int(entity_id), "label": entity_display_label_from_row(entity_row), "reason": normalized_reason, "affected_document_ids": affected_document_ids, "override_count": len(override_keys), "document_links_deleted": document_links_deleted, "resolution_keys_deleted": resolution_keys_deleted, } finally: connection.close() def normalize_document_entity_role(raw_role: object) -> str: normalized = normalize_entity_lookup_text(raw_role).replace(" ", "_") aliases = { "participants": "participant", "recipients": "recipient", "authors": "author", "custodians": "custodian", } role = aliases.get(normalized, normalized) if role not in DOCUMENT_ENTITY_ROLES: raise RetrieverError(f"Unsupported entity role: {raw_role}. Supported roles: {', '.join(sorted(DOCUMENT_ENTITY_ROLES))}.") return role def ensure_document_row(connection: sqlite3.Connection, document_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT * FROM documents WHERE id = ? """, (int(document_id),), ).fetchone() if row is None: raise RetrieverError(f"Unknown document id: {document_id}") return row def next_document_entity_ordinal(connection: sqlite3.Connection, document_id: int, role: str) -> int: row = connection.execute( """ SELECT COALESCE(MAX(ordinal), -1) + 1 AS next_ordinal FROM document_entities WHERE document_id = ? AND role = ? """, (int(document_id), role), ).fetchone() return int(row["next_ordinal"] or 0) if row is not None else 0 def assign_entity( root: Path, *, document_id: int, role: str, entity_id: int, reason: str | None = None, ) -> dict[str, object]: normalized_role = normalize_document_entity_role(role) normalized_reason = normalize_whitespace(str(reason or "")) or None paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="assign-entity"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="assign-entity") raise_if_entity_rebuild_active(connection, root, command_name="assign-entity") ensure_document_row(connection, document_id) entity_row = active_entity_row(connection, entity_id) connection.execute("BEGIN") try: existing_row = connection.execute( """ SELECT * FROM document_entities WHERE document_id = ? AND role = ? AND entity_id = ? ORDER BY id ASC LIMIT 1 """, (int(document_id), normalized_role, int(entity_id)), ).fetchone() evidence = json.dumps( {"source": "manual_assignment", "reason": normalized_reason}, ensure_ascii=True, sort_keys=True, ) if existing_row is not None: connection.execute( """ UPDATE document_entities SET assignment_mode = 'manual', evidence_json = ?, updated_at = ? WHERE id = ? """, (evidence, utc_now(), int(existing_row["id"])), ) created = False else: connection.execute( """ INSERT INTO document_entities ( document_id, entity_id, role, ordinal, assignment_mode, observed_title, evidence_json, created_at, updated_at ) VALUES (?, ?, ?, ?, 'manual', NULL, ?, ?, ?) """, ( int(document_id), int(entity_id), normalized_role, next_document_entity_ordinal(connection, int(document_id), normalized_role), evidence, utc_now(), utc_now(), ), ) created = True refresh_documents_after_entity_graph_change(connection, [int(document_id)]) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "document_id": int(document_id), "role": normalized_role, "entity_id": int(entity_id), "label": entity_display_label_from_row(entity_row), "created": created, "reason": normalized_reason, } finally: connection.close() def document_override_from_link( connection: sqlite3.Connection, *, document_id: int, role: str, source_entity_id: int, override_effect: str, replacement_entity_id: int | None = None, evidence_json: object = None, reason: str | None = None, ) -> int: evidence: dict[str, object] = {} if isinstance(evidence_json, str): try: parsed = json.loads(evidence_json or "{}") if isinstance(parsed, dict): evidence = parsed except json.JSONDecodeError: evidence = {} elif isinstance(evidence_json, dict): evidence = evidence_json normalized_candidate_key = normalize_whitespace(str(evidence.get("normalized_candidate_key") or "")) or None source_hint = normalize_whitespace(str(evidence.get("raw_value") or "")) or None connection.execute( """ INSERT OR IGNORE INTO entity_overrides ( scope_type, scope_id, role, source_entity_id, normalized_candidate_key, replacement_entity_id, override_effect, source_hint, reason, created_at, updated_at ) VALUES ('document', ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( int(document_id), role, int(source_entity_id), normalized_candidate_key, replacement_entity_id, override_effect, source_hint, reason, utc_now(), utc_now(), ), ) return int(connection.execute("SELECT changes()").fetchone()[0] or 0) def unassign_entity( root: Path, *, document_id: int, role: str, entity_id: int, reason: str | None = None, ) -> dict[str, object]: normalized_role = normalize_document_entity_role(role) normalized_reason = normalize_whitespace(str(reason or "")) or None paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="unassign-entity"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="unassign-entity") raise_if_entity_rebuild_active(connection, root, command_name="unassign-entity") ensure_document_row(connection, document_id) active_entity_row(connection, entity_id) link_rows = connection.execute( """ SELECT * FROM document_entities WHERE document_id = ? AND role = ? AND entity_id = ? ORDER BY id ASC """, (int(document_id), normalized_role, int(entity_id)), ).fetchall() if not link_rows: raise RetrieverError( f"Document {document_id} has no {normalized_role} link for entity {entity_id}." ) manual_links_removed = 0 auto_links_removed = 0 overrides_created = 0 connection.execute("BEGIN") try: for link_row in link_rows: if link_row["assignment_mode"] == "manual": connection.execute("DELETE FROM document_entities WHERE id = ?", (int(link_row["id"]),)) manual_links_removed += 1 else: overrides_created += document_override_from_link( connection, document_id=int(document_id), role=normalized_role, source_entity_id=int(entity_id), override_effect="remove", evidence_json=link_row["evidence_json"], reason=normalized_reason, ) connection.execute("DELETE FROM document_entities WHERE id = ?", (int(link_row["id"]),)) auto_links_removed += 1 refresh_documents_after_entity_graph_change(connection, [int(document_id)]) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "document_id": int(document_id), "role": normalized_role, "entity_id": int(entity_id), "manual_links_removed": manual_links_removed, "auto_links_removed": auto_links_removed, "overrides_created": overrides_created, "reason": normalized_reason, } finally: connection.close() def create_split_target_entity( connection: sqlite3.Connection, *, source_row: sqlite3.Row, display_name: str | None, ) -> int: normalized_display_name = normalize_whitespace(str(display_name or "")) or None connection.execute( """ INSERT INTO entities ( entity_type, display_name, display_name_source, entity_origin, canonical_status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( source_row["entity_type"], normalized_display_name, ENTITY_DISPLAY_SOURCE_MANUAL if normalized_display_name else ENTITY_DISPLAY_SOURCE_AUTO, ENTITY_ORIGIN_MANUAL, ENTITY_STATUS_ACTIVE, utc_now(), utc_now(), ), ) return int(connection.execute("SELECT last_insert_rowid()").fetchone()[0]) def selected_identifier_rows_for_split( connection: sqlite3.Connection, *, source_entity_id: int, identifier_ids: list[int] | None, ) -> list[sqlite3.Row]: normalized_ids = list(dict.fromkeys(int(identifier_id) for identifier_id in identifier_ids or [])) if not normalized_ids: return [] placeholders = ", ".join("?" for _ in normalized_ids) rows = connection.execute( f""" SELECT * FROM entity_identifiers WHERE id IN ({placeholders}) AND entity_id = ? ORDER BY id ASC """, [*normalized_ids, int(source_entity_id)], ).fetchall() found_ids = {int(row["id"]) for row in rows} missing_ids = [identifier_id for identifier_id in normalized_ids if identifier_id not in found_ids] if missing_ids: raise RetrieverError( f"Identifier id(s) do not belong to entity {source_entity_id}: {', '.join(str(item) for item in missing_ids)}" ) return rows def move_selected_identifiers_to_target( connection: sqlite3.Connection, *, source_entity_id: int, target_entity_id: int, identifier_rows: list[sqlite3.Row], ) -> dict[str, int]: moved_identifiers = 0 deduped_identifiers = 0 moved_resolution_keys = 0 deleted_resolution_keys = 0 for identifier_row in identifier_rows: source_identifier_id = int(identifier_row["id"]) target_identifier_id = matching_survivor_identifier_id( connection, survivor_entity_id=int(target_entity_id), identifier_row=identifier_row, ) if target_identifier_id is None: connection.execute( """ UPDATE entity_identifiers SET entity_id = ?, source_kind = 'manual', updated_at = ? WHERE id = ? """, (int(target_entity_id), utc_now(), source_identifier_id), ) target_identifier_id = source_identifier_id moved_identifiers += 1 else: deduped_identifiers += 1 for key_row in connection.execute( """ SELECT * FROM entity_resolution_keys WHERE identifier_id = ? OR entity_id = ? AND key_type = ? AND normalized_value = ? AND COALESCE(provider, '') = COALESCE(?, '') AND COALESCE(provider_scope, '') = COALESCE(?, '') AND COALESCE(identifier_name, '') = COALESCE(?, '') AND COALESCE(identifier_scope, '') = COALESCE(?, '') ORDER BY id ASC """, ( source_identifier_id, int(source_entity_id), identifier_row["identifier_type"], identifier_row["normalized_value"], identifier_row["provider"], identifier_row["provider_scope"], identifier_row["identifier_name"], identifier_row["identifier_scope"], ), ).fetchall(): existing_key = resolution_key_matches_row(connection, key_row, exclude_id=int(key_row["id"])) if existing_key is not None: connection.execute("DELETE FROM entity_resolution_keys WHERE id = ?", (int(key_row["id"]),)) deleted_resolution_keys += 1 else: connection.execute( """ UPDATE entity_resolution_keys SET entity_id = ?, identifier_id = ?, updated_at = ? WHERE id = ? """, (int(target_entity_id), target_identifier_id, utc_now(), int(key_row["id"])), ) moved_resolution_keys += 1 if target_identifier_id != source_identifier_id: connection.execute("DELETE FROM entity_identifiers WHERE id = ?", (source_identifier_id,)) return { "moved_identifiers": moved_identifiers, "deduped_identifiers": deduped_identifiers, "moved_resolution_keys": moved_resolution_keys, "deleted_resolution_keys": deleted_resolution_keys, } def selected_document_entity_rows_for_split( connection: sqlite3.Connection, *, source_entity_id: int, document_ids: list[int] | None, roles: list[str] | None, ) -> list[sqlite3.Row]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids or [])) normalized_roles = [normalize_document_entity_role(role) for role in roles or []] if not normalized_document_ids: return [] for document_id in normalized_document_ids: ensure_document_row(connection, document_id) params: list[object] = [int(source_entity_id), *normalized_document_ids] role_clause = "" if normalized_roles: role_clause = f"AND role IN ({', '.join('?' for _ in normalized_roles)})" params.extend(normalized_roles) rows = connection.execute( f""" SELECT * FROM document_entities WHERE entity_id = ? AND document_id IN ({', '.join('?' for _ in normalized_document_ids)}) {role_clause} ORDER BY document_id ASC, role ASC, id ASC """, params, ).fetchall() if not rows: raise RetrieverError("No matching document/entity links found to split.") return rows def split_entity( root: Path, source_entity_id: int, *, target_entity_id: int | None = None, identifier_ids: list[int] | None = None, document_ids: list[int] | None = None, roles: list[str] | None = None, display_name: str | None = None, reason: str | None = None, block_merge: bool = True, ) -> dict[str, object]: if not identifier_ids and not document_ids: raise RetrieverError("split-entity requires at least one --identifier-id or --doc-id.") normalized_reason = normalize_whitespace(str(reason or "")) or None paths = workspace_paths(root) ensure_layout(paths) with workspace_entity_rebuild_session(paths, command_name="split-entity"): connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="split-entity") raise_if_entity_rebuild_active(connection, root, command_name="split-entity") source_row = active_entity_row(connection, source_entity_id) identifier_rows = selected_identifier_rows_for_split( connection, source_entity_id=int(source_entity_id), identifier_ids=identifier_ids, ) document_link_rows = selected_document_entity_rows_for_split( connection, source_entity_id=int(source_entity_id), document_ids=document_ids, roles=roles, ) connection.execute("BEGIN") try: created_target = target_entity_id is None if target_entity_id is None: target_id = create_split_target_entity( connection, source_row=source_row, display_name=display_name, ) else: target_id = int(target_entity_id) active_entity_row(connection, target_id) identifier_counts = move_selected_identifiers_to_target( connection, source_entity_id=int(source_entity_id), target_entity_id=target_id, identifier_rows=identifier_rows, ) moved_links = 0 deduped_links = 0 overrides_created = 0 affected_document_ids: list[int] = [] for link_row in document_link_rows: document_id = int(link_row["document_id"]) role = str(link_row["role"]) affected_document_ids.append(document_id) if link_row["assignment_mode"] == "auto": overrides_created += document_override_from_link( connection, document_id=document_id, role=role, source_entity_id=int(source_entity_id), override_effect="replace", replacement_entity_id=target_id, evidence_json=link_row["evidence_json"], reason=normalized_reason, ) existing_target_link = connection.execute( """ SELECT id FROM document_entities WHERE document_id = ? AND role = ? AND entity_id = ? ORDER BY id ASC LIMIT 1 """, (document_id, role, target_id), ).fetchone() if existing_target_link is not None: connection.execute("DELETE FROM document_entities WHERE id = ?", (int(link_row["id"]),)) deduped_links += 1 else: connection.execute( """ UPDATE document_entities SET entity_id = ?, assignment_mode = 'manual', updated_at = ? WHERE id = ? """, (target_id, utc_now(), int(link_row["id"])), ) moved_links += 1 if block_merge: left, right = entity_pair_key(int(source_entity_id), target_id) connection.execute( """ INSERT OR IGNORE INTO entity_merge_blocks ( left_entity_id, right_entity_id, reason, created_at ) VALUES (?, ?, ?, ?) """, (left, right, normalized_reason or "entity split", utc_now()), ) recompute_entity_caches(connection, int(source_entity_id)) recompute_entity_caches(connection, target_id) refresh_documents_after_entity_graph_change(connection, affected_document_ids) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "source_entity_id": int(source_entity_id), "target_entity_id": target_id, "created_target": created_target, "display_name": normalize_whitespace(str(display_name or "")) or None, "reason": normalized_reason, "block_merge": bool(block_merge), "moved_document_links": moved_links, "deduped_document_links": deduped_links, "overrides_created": overrides_created, "affected_document_ids": sorted(set(affected_document_ids)), **identifier_counts, } finally: connection.close() def list_datasets(root: Path) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) return { "status": "ok", "datasets": list_dataset_summaries(connection, root=root), } finally: connection.close() def create_dataset(root: Path, dataset_name: str) -> dict[str, object]: normalized_name = normalize_whitespace(dataset_name) if not normalized_name: raise RetrieverError("Dataset name cannot be empty.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) existing_rows = find_dataset_rows_by_name(connection, normalized_name) if existing_rows: ids = ", ".join(str(row["id"]) for row in existing_rows) raise RetrieverError( f"Dataset name {normalized_name!r} already exists; use the existing dataset id(s): {ids}." ) connection.execute("BEGIN") try: dataset_id = create_dataset_row(connection, normalized_name) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "dataset": dataset_summary_by_id(connection, dataset_id, root=root), } finally: connection.close() def normalize_dataset_policy_bool(raw_value: object, *, field_name: str) -> bool: if isinstance(raw_value, bool): return raw_value text = normalize_entity_lookup_text(raw_value) if text in {"1", "true", "yes", "y", "on", "enable", "enabled"}: return True if text in {"0", "false", "no", "n", "off", "disable", "disabled"}: return False raise RetrieverError(f"Invalid boolean value for {field_name}: {raw_value!r}") def normalize_dataset_external_id_policy_names(raw_names: list[str] | None) -> list[str]: names: list[str] = [] seen: set[str] = set() for raw_name in raw_names or []: name = normalize_entity_identifier_name(raw_name) if not name: raise RetrieverError(f"Invalid external-id merge policy name: {raw_name!r}") if name in seen: continue seen.add(name) names.append(name) return names def show_dataset_policy( root: Path, *, dataset_id: int | None = None, dataset_name: str | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="set-dataset-policy") dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) return { "status": "ok", "dataset": dataset_summary_by_id(connection, int(dataset_row["id"]), root=root), "merge_policy": dataset_merge_policy_payload_from_row(dataset_row), } finally: connection.close() def set_dataset_policy( root: Path, *, dataset_id: int | None = None, dataset_name: str | None = None, allow_auto_merge: object | None = None, email_auto_merge: object | None = None, handle_auto_merge: object | None = None, phone_auto_merge: object | None = None, name_auto_merge: object | None = None, external_id_auto_merge_names: list[str] | None = None, clear_external_id_auto_merge_names: bool = False, ) -> dict[str, object]: if clear_external_id_auto_merge_names and external_id_auto_merge_names is not None: raise RetrieverError("Use either --external-id-auto-merge-name or --clear-external-id-auto-merge-names, not both.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="set-dataset-policy") raise_if_entity_rebuild_active(connection, root, command_name="set-dataset-policy") dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) if normalize_whitespace(str(dataset_row["source_kind"] or "")).lower() == MANUAL_DATASET_SOURCE_KIND: raise RetrieverError("Dataset merge policy controls apply only to source-backed datasets.") before_policy = dataset_merge_policy_payload_from_row(dataset_row) updates: dict[str, object] = {} for column_name, raw_value in ( ("allow_auto_merge", allow_auto_merge), ("email_auto_merge", email_auto_merge), ("handle_auto_merge", handle_auto_merge), ("phone_auto_merge", phone_auto_merge), ("name_auto_merge", name_auto_merge), ): if raw_value is None: continue updates[column_name] = 1 if normalize_dataset_policy_bool(raw_value, field_name=column_name) else 0 if external_id_auto_merge_names is not None: updates["external_id_auto_merge_names_json"] = json.dumps( normalize_dataset_external_id_policy_names(external_id_auto_merge_names), ensure_ascii=True, sort_keys=True, ) elif clear_external_id_auto_merge_names: updates["external_id_auto_merge_names_json"] = "[]" if not updates: raise RetrieverError("set-dataset-policy requires at least one policy flag.") connection.execute("BEGIN") try: updates["updated_at"] = utc_now() set_clause = ", ".join(f"{quote_identifier(column)} = ?" for column in updates) connection.execute( f""" UPDATE datasets SET {set_clause} WHERE id = ? """, [*updates.values(), int(dataset_row["id"])], ) connection.commit() except Exception: connection.rollback() raise updated_row = get_dataset_row_by_id(connection, int(dataset_row["id"])) assert updated_row is not None after_policy = dataset_merge_policy_payload_from_row(updated_row) return { "status": "ok", "dataset": dataset_summary_by_id(connection, int(updated_row["id"]), root=root), "before_merge_policy": before_policy, "merge_policy": after_policy, "changed": before_policy != after_policy, "rebuild_recommended": before_policy != after_policy, "rebuild_command": f"rebuild-entities {shlex.quote(str(root))}", } finally: connection.close() def add_to_dataset( root: Path, document_ids: list[int], *, dataset_id: int | None = None, dataset_name: str | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) connection.execute("BEGIN") try: result = add_documents_to_dataset( connection, dataset_id=int(dataset_row["id"]), document_ids=document_ids, ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "requested_document_ids": sorted(dict.fromkeys(int(document_id) for document_id in document_ids)), **result, "dataset": dataset_summary_by_id(connection, int(dataset_row["id"]), root=root), } finally: connection.close() def remove_from_dataset( root: Path, document_ids: list[int], *, dataset_id: int | None = None, dataset_name: str | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) connection.execute("BEGIN") try: result = remove_documents_from_dataset( connection, dataset_id=int(dataset_row["id"]), document_ids=document_ids, ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "requested_document_ids": sorted(dict.fromkeys(int(document_id) for document_id in document_ids)), **result, "dataset": dataset_summary_by_id(connection, int(dataset_row["id"]), root=root), } finally: connection.close() def delete_dataset( root: Path, *, dataset_id: int | None = None, dataset_name: str | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) connection.execute("BEGIN") try: result = delete_dataset_row(connection, int(dataset_row["id"]), root=root) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", **result, } finally: connection.close() def rename_dataset(root: Path, old_name: str, new_name: str) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) dataset_row = resolve_dataset_row(connection, dataset_name=old_name) connection.execute("BEGIN") try: renamed_summary = rename_dataset_row(connection, int(dataset_row["id"]), new_name, root=root) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "dataset": renamed_summary, } finally: connection.close() def list_runs(root: Path) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) return { "status": "ok", "runs": list_run_summaries(connection), } finally: connection.close() def get_run(root: Path, run_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) return { "status": "ok", "run": run_summary_by_id(connection, run_id), } finally: connection.close() def create_run( root: Path, *, job_version_id: int | None = None, raw_job_name: str | None = None, job_version_number: int | None = None, dataset_names: list[str] | None = None, document_ids: list[int] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, activation_policy: str = "manual", family_mode: str = "exact", seed_limit: int | None = None, ) -> dict[str, object]: normalized_job_name = ( sanitize_processing_identifier(raw_job_name, label="Job name", prefix="job") if raw_job_name is not None else None ) normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) normalized_family_mode = normalize_run_family_mode(family_mode) normalized_activation_policy = normalize_run_activation_policy(activation_policy) if normalized_document_ids and (query.strip() or raw_bates or raw_filters or dataset_names or from_run_id is not None or select_from_scope): raise RetrieverError("create-run accepts either --doc-id selectors or scope/query selectors, not both.") if seed_limit is not None and seed_limit < 1: raise RetrieverError("Run limit must be >= 1.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) if normalized_document_ids: selector = {"document_ids": normalized_document_ids} preferred_from_run_id = None else: selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if not scope_run_selector_has_inputs(selector): raise RetrieverError("Run selector must include at least one inclusion input.") preferred_from_run_id = preferred_scope_selector_from_run_id(selector) job_version_row = require_job_version_row( connection, job_version_id=job_version_id, job_name=normalized_job_name, version=job_version_number, ) job_row = connection.execute( "SELECT * FROM jobs WHERE id = ?", (job_version_row["job_id"],), ).fetchone() assert job_row is not None job_kind = normalize_job_kind(str(job_row["job_kind"])) if normalized_activation_policy != "manual" and job_kind not in REVISION_PRODUCING_JOB_KINDS: raise RetrieverError( f"Run activation policy '{normalized_activation_policy}' is only supported for " f"revision-producing jobs ({', '.join(sorted(REVISION_PRODUCING_JOB_KINDS))}); " f"job '{job_row['job_name']}' is kind '{job_kind}'." ) if normalized_document_ids: selected_document_rows = fetch_visible_document_rows_by_ids(connection, normalized_document_ids) seed_document_ids = [int(row["id"]) for row in selected_document_rows] if seed_limit is not None: seed_document_ids = seed_document_ids[:seed_limit] reasons_by_document_id = { document_id: { "direct_reasons": [{"type": "document_id", "document_id": document_id}], "family_seed_document_ids": [], } for document_id in seed_document_ids } if normalized_family_mode == "with_family": final_document_ids = expand_seed_documents_with_family(connection, seed_document_ids, reasons_by_document_id) else: final_document_ids = list(seed_document_ids) if final_document_ids: document_rows = connection.execute( f""" SELECT * FROM documents WHERE id IN ({', '.join('?' for _ in final_document_ids)}) ORDER BY id ASC """, final_document_ids, ).fetchall() document_row_by_id = {int(row["id"]): row for row in document_rows} else: document_row_by_id = {} snapshot_rows = [] for ordinal, document_id in enumerate(final_document_ids): document_row = document_row_by_id.get(int(document_id)) if document_row is None: continue pinned_input = compute_document_input_reference_for_job_version( connection, root=root, document_row=document_row, job_row=job_row, job_version_row=job_version_row, frozen_input_revision_id=None, frozen_content_hash=None, ) snapshot_rows.append( { "document_id": int(document_id), "ordinal": ordinal, "inclusion_reason": reasons_by_document_id.get( int(document_id), {"direct_reasons": [], "family_seed_document_ids": []}, ), "pinned_input_revision_id": pinned_input["pinned_input_revision_id"], "pinned_input_identity": pinned_input["pinned_input_identity"], "pinned_content_hash": pinned_input["pinned_content_hash"], } ) else: snapshot_rows = plan_scope_run_snapshot_rows( connection, root=root, job_row=job_row, job_version_row=job_version_row, selector=selector, family_mode=normalized_family_mode, seed_limit=seed_limit, ) connection.execute("BEGIN") try: run_id = create_run_row( connection, job_version_id=int(job_version_row["id"]), selector=selector, exclude_selector={}, activation_policy=normalized_activation_policy, family_mode=normalized_family_mode, seed_limit=seed_limit, from_run_id=preferred_from_run_id, status="planned", ) replace_run_snapshot_documents( connection, run_id=run_id, snapshot_rows=snapshot_rows, ) if normalize_job_kind(str(job_row["job_kind"])) == "ocr": materialize_run_items_for_run(connection, paths, root, run_id) refresh_run_progress(connection, run_id) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "run": run_summary_by_id(connection, run_id), } finally: connection.close() def list_text_revisions(root: Path, *, document_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) document_row = connection.execute( """ SELECT id FROM documents WHERE id = ? """, (document_id,), ).fetchone() if document_row is None: raise RetrieverError(f"Unknown document id: {document_id}") return { "status": "ok", "document_id": int(document_id), "text_revisions": list_text_revision_summaries_for_document(connection, int(document_id)), } finally: connection.close() def activate_text_revision( root: Path, *, document_id: int, text_revision_id: int, activation_policy: str = "manual", ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN") try: payload = activate_text_revision_for_document( connection, paths, document_id=document_id, text_revision_id=text_revision_id, activation_policy=activation_policy, ) connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def list_results( root: Path, *, run_id: int | None = None, document_id: int | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) return { "status": "ok", "results": list_result_summaries(connection, run_id=run_id, document_id=document_id), } finally: connection.close() def claim_run_items( root: Path, *, run_id: int, claimed_by: str, limit: int = DEFAULT_RUN_ITEM_CLAIM_BATCH_SIZE, stale_after_seconds: int | None = None, launch_mode: str = "inline", worker_task_id: str | None = None, max_batches: int | None = None, ) -> dict[str, object]: if limit < 1: raise RetrieverError("Claim limit must be >= 1.") effective_stale_after_seconds = ( int(stale_after_seconds) if stale_after_seconds is not None else default_run_item_claim_stale_seconds_for_launch_mode(launch_mode) ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN IMMEDIATE") try: materialize_run_items_for_run(connection, paths, root, run_id) ensure_run_worker_row( connection, run_id=run_id, claimed_by=claimed_by, launch_mode=launch_mode, worker_task_id=worker_task_id, max_batches=max_batches, ) reused_count = reuse_active_results_for_run(connection, run_id) claimed_rows = claim_run_item_rows( connection, run_id=run_id, claimed_by=claimed_by, limit=limit, stale_after_seconds=effective_stale_after_seconds, ) if claimed_rows: update_run_worker_row( connection, run_id=run_id, claimed_by=normalize_whitespace(claimed_by), heartbeat=True, increment_batches_prepared=True, ) refresh_run_progress(connection, run_id) payload = { "status": "ok", "run": run_status_by_id(connection, run_id), "claimed_by": normalize_whitespace(claimed_by), "stale_after_seconds": effective_stale_after_seconds, "reused_count": reused_count, "run_items": [run_item_row_to_payload(row) for row in claimed_rows], } connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def prepare_run_batch( root: Path, *, run_id: int, claimed_by: str, limit: int | None = None, stale_after_seconds: int | None = None, launch_mode: str = "inline", worker_task_id: str | None = None, max_batches: int | None = None, budget_seconds: int | None = None, ) -> dict[str, object]: if limit is not None and limit < 1: raise RetrieverError("Claim limit must be >= 1.") budget = normalize_resumable_step_budget(budget_seconds) if budget_seconds is not None else None effective_stale_after_seconds = ( int(stale_after_seconds) if stale_after_seconds is not None else default_run_item_claim_stale_seconds_for_launch_mode(launch_mode) ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN IMMEDIATE") try: materialize_run_items_for_run(connection, paths, root, run_id) ensure_run_worker_row( connection, run_id=run_id, claimed_by=claimed_by, launch_mode=launch_mode, worker_task_id=worker_task_id, max_batches=max_batches, ) reused_count = reuse_active_results_for_run(connection, run_id) initial_run_payload = run_status_by_id( connection, run_id, claimed_by=claimed_by, ) initial_worker_payload = dict(initial_run_payload.get("worker") or {}) effective_limit = limit if limit is not None else int(initial_worker_payload["recommended_batch_size"]) if budget is not None and budget < RUN_JOB_MIN_SECONDS_TO_CLAIM and initial_worker_payload["next_action"] == "claim": effective_limit = 0 claimed_rows: list[sqlite3.Row] = [] batch_payloads: list[dict[str, object]] = [] if initial_worker_payload["next_action"] == "process_batch": claimed_rows = claimed_run_item_rows_for_worker( connection, run_id=run_id, claimed_by=claimed_by, ) if claimed_rows: heartbeat_claimed_run_items(connection, run_id=run_id, claimed_by=claimed_by) batch_payloads = [ { "run_item": run_item_row_to_payload(row), "context": build_run_item_context_payload(connection, paths, root, row), } for row in claimed_rows ] elif initial_worker_payload["next_action"] == "claim" and effective_limit > 0: claimed_rows = claim_run_item_rows( connection, run_id=run_id, claimed_by=claimed_by, limit=effective_limit, stale_after_seconds=effective_stale_after_seconds, ) batch_payloads = [ { "run_item": run_item_row_to_payload(row), "context": build_run_item_context_payload(connection, paths, root, row), } for row in claimed_rows ] if batch_payloads: update_run_worker_row( connection, run_id=run_id, claimed_by=normalize_whitespace(claimed_by), heartbeat=True, increment_batches_prepared=True, ) current_run_payload = run_status_by_id( connection, run_id, claimed_by=claimed_by, ) worker_payload = dict(current_run_payload.get("worker") or {}) if budget is not None: worker_payload["budget_seconds"] = budget worker_payload["minimum_seconds_to_claim"] = RUN_JOB_MIN_SECONDS_TO_CLAIM if budget < RUN_JOB_MIN_SECONDS_TO_CLAIM and worker_payload["next_action"] == "claim": worker_payload["next_action"] = "stop" worker_payload["stop_reason"] = "budget_exhausted" if batch_payloads: worker_payload["next_action"] = "process_batch" worker_payload["stop_reason"] = None elif worker_payload["next_action"] == "claim": worker_payload["next_action"] = "stop" worker_payload["stop_reason"] = "no_claimable_items" worker_payload["prepared_batch_size"] = len(batch_payloads) current_run_payload["next_recommended_commands"] = run_job_next_recommended_commands( root, run_payload=current_run_payload, budget_seconds=budget if budget is not None else DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, claimed_by=claimed_by, ) payload = { "status": "ok", "run": current_run_payload, "worker": worker_payload, "claimed_by": normalize_whitespace(claimed_by), "requested_limit": effective_limit, "stale_after_seconds": effective_stale_after_seconds, "reused_count": reused_count, "batch": batch_payloads, "worker_record": run_worker_row_to_payload( find_run_worker_row(connection, run_id=run_id, claimed_by=normalize_whitespace(claimed_by)) ), } connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def get_run_item_context(root: Path, *, run_item_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) run_item_row = require_run_item_row_by_id(connection, run_item_id) return { "status": "ok", "context": build_run_item_context_payload(connection, paths, root, run_item_row), } finally: connection.close() def heartbeat_run_items( root: Path, *, run_id: int, claimed_by: str, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN") try: updated_count = heartbeat_claimed_run_items(connection, run_id=run_id, claimed_by=claimed_by) payload = { "status": "ok", "updated_count": updated_count, "run": run_status_by_id(connection, run_id), } connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def complete_run_item( root: Path, *, run_item_id: int, claimed_by: str, page_text: str | None = None, raw_output_json: str | None = None, normalized_output_json: str | None = None, output_values_json: str | None = None, created_text_revision_json: str | None = None, provider_metadata_json: str | None = None, provider_request_id: str | None = None, input_tokens: int | None = None, output_tokens: int | None = None, cost_cents: int | None = None, latency_ms: int | None = None, ) -> dict[str, object]: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN") try: run_item_row = require_run_item_row_by_id(connection, run_item_id) if str(run_item_row["status"] or "") == "completed": result_payload = ( result_summary_by_id(connection, int(run_item_row["result_id"])) if run_item_row["result_id"] is not None else None ) ocr_page_output_payload = None image_description_page_output_payload = None if str(run_item_row["item_kind"] or "") == "page": run_row = require_run_row_by_id(connection, int(run_item_row["run_id"])) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( "SELECT * FROM jobs WHERE id = ?", (job_version_row["job_id"],), ).fetchone() assert job_row is not None page_job_kind = normalize_job_kind(str(job_row["job_kind"])) if page_job_kind == "ocr": existing_page_output_row = find_ocr_page_output_row(connection, run_item_id=run_item_id) if existing_page_output_row is not None: ocr_page_output_payload = ocr_page_output_row_to_payload(existing_page_output_row) elif page_job_kind == "image_description": existing_page_output_row = find_image_description_page_output_row( connection, run_item_id=run_item_id, ) if existing_page_output_row is not None: image_description_page_output_payload = image_description_page_output_row_to_payload( existing_page_output_row ) payload = { "status": "ok", "idempotent": True, "run_item": run_item_row_to_payload(run_item_row), "result": result_payload, "ocr_page_output": ocr_page_output_payload, "image_description_page_output": image_description_page_output_payload, "run": run_status_by_id(connection, int(run_item_row["run_id"])), } connection.commit() return payload if str(run_item_row["status"] or "") == "failed": raise RetrieverError(f"Run item {run_item_id} is already failed; reclaim it before completing it.") current_claimed_by = normalize_whitespace(str(run_item_row["claimed_by"] or "")) if current_claimed_by and current_claimed_by != normalized_claimed_by: raise RetrieverError( f"Run item {run_item_id} is claimed by {current_claimed_by!r}, not {normalized_claimed_by!r}." ) if str(run_item_row["status"] or "") != "running": raise RetrieverError(f"Run item {run_item_id} must be running before it can be completed.") run_row = require_run_row_by_id(connection, int(run_item_row["run_id"])) job_version_row = require_job_version_row_by_id(connection, int(run_row["job_version_id"])) job_row = connection.execute( "SELECT * FROM jobs WHERE id = ?", (job_version_row["job_id"],), ).fetchone() assert job_row is not None job_kind = normalize_job_kind(str(job_row["job_kind"])) snapshot_row = None if run_item_row["run_snapshot_document_id"] is not None: snapshot_row = connection.execute( "SELECT * FROM run_snapshot_documents WHERE id = ?", (run_item_row["run_snapshot_document_id"],), ).fetchone() job_output_rows = connection.execute( """ SELECT * FROM job_outputs WHERE job_id = ? ORDER BY ordinal ASC, output_name ASC, id ASC """, (job_row["id"],), ).fetchall() raw_output = parse_json_argument(raw_output_json, label="Raw output", default=None) normalized_output = parse_json_argument( normalized_output_json, label="Normalized output", default=raw_output, ) output_values_default = normalized_output if isinstance(normalized_output, dict) else {} output_values = parse_json_object_argument( output_values_json, label="Output values", default=output_values_default if isinstance(output_values_default, dict) else {}, ) provider_metadata = dict( parse_json_object_argument( provider_metadata_json, label="Provider metadata", default={}, ) ) created_text_revision_payload = ( parse_json_object_argument( created_text_revision_json, label="Created text revision", default={}, ) if created_text_revision_json is not None else None ) if job_kind == "structured_extraction": schema_issues = validate_processing_response_output( job_version_row, job_output_rows, output_values, ) provider_metadata["validation"] = { "kind": "response_schema", "status": "failed" if schema_issues else "ok", "issue_count": len(schema_issues), "issues": list(schema_issues), } if schema_issues: sample = "; ".join(schema_issues[:3]) suffix = f" (+{len(schema_issues) - 3} more)" if len(schema_issues) > 3 else "" raise RetrieverError(f"Structured output did not match response_schema: {sample}{suffix}") created_text_revision_id = None if job_kind in {"ocr", "image_description"} and str(run_item_row["item_kind"] or "") == "page": resolved_page_text = page_text if page_text is not None else None if resolved_page_text is None: normalized_candidate = normalized_output if isinstance(normalized_output, str) else None raw_candidate = raw_output if isinstance(raw_output, str) else None resolved_page_text = normalized_candidate or raw_candidate if resolved_page_text is None: raise RetrieverError( "Page completion requires --page-text or a raw/normalized string payload." ) page_raw_output = raw_output if raw_output is not None else {"page_text": resolved_page_text} page_normalized_output = ( normalized_output if normalized_output is not None else {"page_text": resolved_page_text} ) page_output_payload_key = "ocr_page_output" if job_kind == "ocr": upsert_ocr_page_output_row( connection, run_item_id=run_item_id, run_id=int(run_item_row["run_id"]), document_id=int(run_item_row["document_id"]), page_number=int(run_item_row["page_number"] or 0), text_content=str(resolved_page_text), raw_output=page_raw_output, normalized_output=page_normalized_output, provider_metadata=provider_metadata, ) page_output_payload = ocr_page_output_row_to_payload( find_ocr_page_output_row(connection, run_item_id=run_item_id) # type: ignore[arg-type] ) else: page_output_payload_key = "image_description_page_output" upsert_image_description_page_output_row( connection, run_item_id=run_item_id, run_id=int(run_item_row["run_id"]), document_id=int(run_item_row["document_id"]), page_number=int(run_item_row["page_number"] or 0), text_content=str(resolved_page_text), raw_output=page_raw_output, normalized_output=page_normalized_output, provider_metadata=provider_metadata, ) page_output_payload = image_description_page_output_row_to_payload( find_image_description_page_output_row(connection, run_item_id=run_item_id) # type: ignore[arg-type] ) create_attempt_row( connection, run_item_id=run_item_id, provider_request_id=provider_request_id, input_tokens=input_tokens, output_tokens=output_tokens, cost_cents=cost_cents, latency_ms=latency_ms, provider_metadata=provider_metadata or {"executor": "cowork_agent"}, error_summary=None, ) completion_time = utc_now() update_run_item_row( connection, run_item_id=run_item_id, status="completed", result_id=None, last_error=None, claimed_by=normalized_claimed_by, claimed_at=str(run_item_row["claimed_at"] or completion_time), last_heartbeat_at=completion_time, completed_at=completion_time, increment_attempt_count=True, ) update_run_worker_row( connection, run_id=int(run_item_row["run_id"]), claimed_by=normalized_claimed_by, heartbeat=True, increment_items_completed=1, ) payload = { "status": "ok", "idempotent": False, "run_item": run_item_row_to_payload(require_run_item_row_by_id(connection, run_item_id)), "run": run_status_by_id(connection, int(run_item_row["run_id"])), } payload[page_output_payload_key] = page_output_payload connection.commit() return payload if created_text_revision_payload: text_content = str(created_text_revision_payload.get("text_content") or "") if not text_content: raise RetrieverError("Created text revision payload must include text_content.") revision_kind = str(created_text_revision_payload.get("revision_kind") or job_kind) revision_language = ( normalize_whitespace(str(created_text_revision_payload.get("language") or "")).lower() or None ) source_text = None if snapshot_row is not None and snapshot_row["pinned_input_revision_id"] is not None: source_revision_row = require_text_revision_row_by_id( connection, int(snapshot_row["pinned_input_revision_id"]), ) source_text = read_text_revision_body(paths, source_revision_row["storage_rel_path"]) or "" quality_summary = None if revision_kind == "translation" or job_kind == "translation": job_parameters = decode_json_text(job_version_row["parameters_json"], default={}) or {} if not isinstance(job_parameters, dict): job_parameters = {} target_language = revision_language if target_language is None: for raw_value in ( (normalized_output.get("target_language") if isinstance(normalized_output, dict) else None), (raw_output.get("target_language") if isinstance(raw_output, dict) else None), job_parameters.get("target_language"), job_parameters.get("target_lang"), job_parameters.get("language"), ): normalized_value = normalize_whitespace(str(raw_value or "")).lower() if normalized_value: target_language = normalized_value break validation_summary = ( translation_validation_summary( source_text or "", text_content, target_language=target_language, ) if source_text is not None else None ) if validation_summary is not None: provider_metadata["validation"] = validation_summary blocking_issues = list(validation_summary.get("blocking_issues") or []) if blocking_issues: sample = "; ".join(blocking_issues[:3]) suffix = f" (+{len(blocking_issues) - 3} more)" if len(blocking_issues) > 3 else "" raise RetrieverError( f"Translation literal-preservation validation failed: {sample}{suffix}" ) quality_summary = text_quality_summary( text_content, revision_kind="translation", source_text=source_text, validation_summary=validation_summary if isinstance(validation_summary, dict) else None, ) else: quality_summary = text_quality_summary( text_content, revision_kind=revision_kind, source_text=source_text, ) provider_metadata["quality"] = quality_summary created_quality_score = created_text_revision_payload.get("quality_score") created_text_revision_id = create_text_revision_row( connection, paths, document_id=int(run_item_row["document_id"]), revision_kind=revision_kind, text_content=text_content, language=revision_language, parent_revision_id=( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row is not None and snapshot_row["pinned_input_revision_id"] is not None else None ), created_by_job_version_id=int(job_version_row["id"]), quality_score=( float(created_quality_score) if created_quality_score is not None else float(quality_summary["quality_score"]) ), provider_metadata=provider_metadata, ) elif job_kind == "translation": raise RetrieverError("Translation run items must include a created text revision payload.") if raw_output is None: if created_text_revision_id is not None: raw_output = { "created_text_revision_id": created_text_revision_id, "job_kind": job_kind, } else: raw_output = output_values if normalized_output is None: normalized_output = output_values if output_values else raw_output result_id, created = create_result_row( connection, run_id=int(run_item_row["run_id"]), document_id=int(run_item_row["document_id"]), job_version_id=int(job_version_row["id"]), input_revision_id=( int(snapshot_row["pinned_input_revision_id"]) if snapshot_row is not None and snapshot_row["pinned_input_revision_id"] is not None else None ), input_identity=str(run_item_row["input_identity"]), raw_output=raw_output, normalized_output=normalized_output, created_text_revision_id=created_text_revision_id, provider_metadata=provider_metadata, ) if created and job_output_rows: upsert_result_output_rows( connection, result_id=result_id, job_output_rows=job_output_rows, output_values_by_name=output_values, ) result_row = connection.execute( """ SELECT * FROM results WHERE id = ? """, (result_id,), ).fetchone() assert result_row is not None activation_payload = maybe_activate_created_text_revision( connection, paths, run_row=require_run_row_by_id(connection, int(run_item_row["run_id"])), job_version_row=job_version_row, document_id=int(run_item_row["document_id"]), result_id=result_id, text_revision_id=( int(result_row["created_text_revision_id"]) if result_row["created_text_revision_id"] is not None else None ), ) create_attempt_row( connection, run_item_id=run_item_id, provider_request_id=provider_request_id, input_tokens=input_tokens, output_tokens=output_tokens, cost_cents=cost_cents, latency_ms=latency_ms, provider_metadata=provider_metadata or {"executor": "cowork_agent"}, error_summary=None, ) completion_time = utc_now() update_run_item_row( connection, run_item_id=run_item_id, status="completed", result_id=result_id, last_error=None, claimed_by=normalized_claimed_by, claimed_at=str(run_item_row["claimed_at"] or completion_time), last_heartbeat_at=completion_time, completed_at=completion_time, increment_attempt_count=True, ) update_run_worker_row( connection, run_id=int(run_item_row["run_id"]), claimed_by=normalized_claimed_by, heartbeat=True, increment_items_completed=1, ) payload = { "status": "ok", "idempotent": False, "run_item": run_item_row_to_payload(require_run_item_row_by_id(connection, run_item_id)), "result": result_summary_by_id(connection, result_id), "run": run_status_by_id(connection, int(run_item_row["run_id"])), } if activation_payload is not None: payload["activation"] = activation_payload connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def fail_run_item( root: Path, *, run_item_id: int, claimed_by: str, error_summary: str, provider_metadata_json: str | None = None, provider_request_id: str | None = None, input_tokens: int | None = None, output_tokens: int | None = None, cost_cents: int | None = None, latency_ms: int | None = None, ) -> dict[str, object]: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") normalized_error = normalize_whitespace(error_summary) if not normalized_error: raise RetrieverError("error cannot be empty.") provider_metadata = parse_json_object_argument( provider_metadata_json, label="Provider metadata", default={}, ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN") try: run_item_row = require_run_item_row_by_id(connection, run_item_id) if str(run_item_row["status"] or "") == "failed": payload = { "status": "ok", "idempotent": True, "run_item": run_item_row_to_payload(run_item_row), "run": run_status_by_id(connection, int(run_item_row["run_id"])), } connection.commit() return payload if str(run_item_row["status"] or "") == "completed": raise RetrieverError(f"Run item {run_item_id} is already completed and cannot be failed.") current_claimed_by = normalize_whitespace(str(run_item_row["claimed_by"] or "")) if current_claimed_by and current_claimed_by != normalized_claimed_by: raise RetrieverError( f"Run item {run_item_id} is claimed by {current_claimed_by!r}, not {normalized_claimed_by!r}." ) if str(run_item_row["status"] or "") != "running": raise RetrieverError(f"Run item {run_item_id} must be running before it can be failed.") create_attempt_row( connection, run_item_id=run_item_id, provider_request_id=provider_request_id, input_tokens=input_tokens, output_tokens=output_tokens, cost_cents=cost_cents, latency_ms=latency_ms, provider_metadata=provider_metadata or {"executor": "cowork_agent"}, error_summary=normalized_error, ) completion_time = utc_now() update_run_item_row( connection, run_item_id=run_item_id, status="failed", result_id=None, last_error=normalized_error, claimed_by=normalized_claimed_by, claimed_at=str(run_item_row["claimed_at"] or completion_time), last_heartbeat_at=completion_time, completed_at=completion_time, increment_attempt_count=True, ) update_run_worker_row( connection, run_id=int(run_item_row["run_id"]), claimed_by=normalized_claimed_by, heartbeat=True, increment_items_failed=1, last_error=normalized_error, ) payload = { "status": "ok", "idempotent": False, "run_item": run_item_row_to_payload(require_run_item_row_by_id(connection, run_item_id)), "run": run_status_by_id(connection, int(run_item_row["run_id"])), } connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def run_job_default_claimed_by(run_id: int) -> str: return f"cowork-run-{int(run_id)}" def run_job_effective_claimed_by(run_id: int, claimed_by: str | None = None) -> str: return normalize_whitespace(claimed_by or run_job_default_claimed_by(run_id)) def run_job_next_recommended_commands( root: Path, *, run_payload: dict[str, object], budget_seconds: int, claimed_by: str | None = None, ) -> list[str]: run_id = int(run_payload["id"]) root_arg = shlex.quote(str(root)) run_id_arg = str(run_id) budget_arg = str(int(budget_seconds)) effective_claimed_by = run_job_effective_claimed_by(run_id, claimed_by) claimed_by_arg = shlex.quote(effective_claimed_by) if str(run_payload.get("status") or "") in {"completed", "failed", "canceled"}: return [f"run-status {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}"] worker_payload = dict(run_payload.get("worker") or {}) next_action = str(worker_payload.get("next_action") or "") handoff_claimed_by_hint = ( normalize_whitespace(str(worker_payload.get("handoff_claimed_by_hint") or "")) or None ) commands: list[str] = [] if next_action == "finalize_ocr": commands.append(f"finalize-ocr-run {root_arg} --run-id {run_id_arg}") elif next_action == "finalize_image_description": commands.append(f"finalize-image-description-run {root_arg} --run-id {run_id_arg}") elif next_action == "handoff": continuation_claimed_by = handoff_claimed_by_hint or f"{effective_claimed_by}-handoff" continuation_arg = shlex.quote(continuation_claimed_by) commands.extend( [ "run-job-step " f"{root_arg} --run-id {run_id_arg} --claimed-by {continuation_arg} --budget-seconds {budget_arg}", "prepare-run-batch " f"{root_arg} --run-id {run_id_arg} --claimed-by {continuation_arg} " f"--budget-seconds {budget_arg} --stale-seconds {DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS}", ] ) elif next_action in {"claim", "process_batch", "stop"}: run_item_counts = dict(run_payload.get("run_item_counts") or {}) if ( int(worker_payload.get("outstanding_items", 0) or 0) > 0 or int(run_item_counts.get("pending", 0) or 0) > 0 or int(run_item_counts.get("running", 0) or 0) > 0 ): commands.append( "prepare-run-batch " f"{root_arg} --run-id {run_id_arg} --claimed-by {claimed_by_arg} " f"--budget-seconds {budget_arg} --stale-seconds {DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS}" ) if commands and next_action != "handoff": commands.insert( 0, "run-job-step " f"{root_arg} --run-id {run_id_arg} --claimed-by {claimed_by_arg} --budget-seconds {budget_arg}", ) commands.append(f"run-status {root_arg} --run-id {run_id_arg} --budget-seconds {budget_arg}") return commands def run_status( root: Path, *, run_id: int, budget_seconds: int | None = None, claimed_by: str | None = None, ) -> dict[str, object]: budget = ( normalize_resumable_step_budget(budget_seconds) if budget_seconds is not None else DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) run_payload = run_status_by_id(connection, run_id, claimed_by=claimed_by) run_payload["next_recommended_commands"] = run_job_next_recommended_commands( root, run_payload=run_payload, budget_seconds=budget, claimed_by=claimed_by, ) return { "status": "ok", "run": run_payload, } finally: connection.close() def run_job_step( root: Path, *, run_id: int, claimed_by: str | None = None, budget_seconds: int | None = None, limit: int | None = None, launch_mode: str = "inline", worker_task_id: str | None = None, max_batches: int | None = None, stale_after_seconds: int | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) normalized_claimed_by = run_job_effective_claimed_by(run_id, claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") status_payload = run_status( root, run_id=run_id, budget_seconds=budget, claimed_by=normalized_claimed_by, ) run_payload = dict(status_payload["run"]) worker_payload = dict(run_payload.get("worker") or {}) next_action = str(worker_payload.get("next_action") or "") if str(run_payload.get("status") or "") in {"completed", "failed", "canceled"}: return { "status": "ok", "step": "run-job", "run_id": run_id, "claimed_by": normalized_claimed_by, "budget_seconds": budget, "executed": False, "executed_step": None, "reason": "run_terminal", "batch": [], "step_result": None, "run": run_payload, "more_work_remaining": False, "next_recommended_commands": run_payload.get("next_recommended_commands", []), } if budget < RUN_JOB_MIN_SECONDS_TO_CLAIM and next_action == "claim": run_payload["next_recommended_commands"] = run_job_next_recommended_commands( root, run_payload=run_payload, budget_seconds=budget, claimed_by=normalized_claimed_by, ) return { "status": "ok", "step": "run-job", "run_id": run_id, "claimed_by": normalized_claimed_by, "budget_seconds": budget, "executed": False, "executed_step": None, "reason": "budget_exhausted", "batch": [], "step_result": None, "run": run_payload, "more_work_remaining": True, "next_recommended_commands": run_payload["next_recommended_commands"], } executed_step: str | None = None step_result: dict[str, object] | None = None if next_action == "finalize_ocr": executed_step = "finalize_ocr" step_result = finalize_ocr_run(root, run_id=run_id) elif next_action == "finalize_image_description": executed_step = "finalize_image_description" step_result = finalize_image_description_run(root, run_id=run_id) elif next_action in {"claim", "process_batch"}: executed_step = "prepare_run_batch" step_result = prepare_run_batch( root, run_id=run_id, claimed_by=normalized_claimed_by, limit=limit, stale_after_seconds=( stale_after_seconds if stale_after_seconds is not None else DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS ), launch_mode=launch_mode, worker_task_id=worker_task_id, max_batches=max_batches, budget_seconds=budget, ) else: run_payload["next_recommended_commands"] = run_job_next_recommended_commands( root, run_payload=run_payload, budget_seconds=budget, claimed_by=normalized_claimed_by, ) handoff_claimed_by_hint = ( normalize_whitespace(str(worker_payload.get("handoff_claimed_by_hint") or "")) or None ) return { "status": "ok", "step": "run-job", "run_id": run_id, "claimed_by": normalized_claimed_by, "budget_seconds": budget, "executed": False, "executed_step": None, "reason": "handoff_required" if next_action == "handoff" else (next_action or "stop"), "batch": [], "step_result": None, "run": run_payload, "handoff_claimed_by_hint": handoff_claimed_by_hint, "more_work_remaining": bool(run_payload.get("status") not in {"completed", "failed", "canceled"}), "next_recommended_commands": run_payload["next_recommended_commands"], } updated_run_payload = dict( (step_result or {}).get("run") or run_status( root, run_id=run_id, budget_seconds=budget, claimed_by=normalized_claimed_by, )["run"] ) updated_run_payload["next_recommended_commands"] = run_job_next_recommended_commands( root, run_payload=updated_run_payload, budget_seconds=budget, claimed_by=normalized_claimed_by, ) batch_payload = list((step_result or {}).get("batch") or []) return { "status": "ok", "step": "run-job", "run_id": run_id, "claimed_by": normalized_claimed_by, "budget_seconds": budget, "executed": True, "executed_step": executed_step, "reason": ( "batch_ready" if batch_payload else ( "run_terminal" if str(updated_run_payload.get("status") or "") in {"completed", "failed", "canceled"} else "step_complete" ) ), "batch": batch_payload, "worker": (step_result or {}).get("worker"), "step_result": step_result, "run": updated_run_payload, "more_work_remaining": bool( str(updated_run_payload.get("status") or "") not in {"completed", "failed", "canceled"} ), "next_recommended_commands": updated_run_payload["next_recommended_commands"], } def cancel_run(root: Path, *, run_id: int, force: bool = False) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN IMMEDIATE") try: run_row = require_run_row_by_id(connection, run_id) already_canceled = str(run_row["status"] or "") == "canceled" or run_row["canceled_at"] is not None materialize_run_items_for_run(connection, paths, root, run_id) canceled_at = str(run_row["canceled_at"] or utc_now()) connection.execute( """ UPDATE runs SET status = 'canceled', canceled_at = COALESCE(canceled_at, ?) WHERE id = ? """, (canceled_at, run_id), ) skipped_count = cancel_pending_run_items(connection, run_id=run_id) force_stop_task_ids = request_run_worker_cancellation(connection, run_id=run_id, force=force) payload = { "status": "ok", "idempotent": already_canceled, "canceled_pending_items": skipped_count, "force_stop_requested": force, "force_stop_task_ids": force_stop_task_ids, "run": run_status_by_id(connection, run_id), } connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def finish_run_worker( root: Path, *, run_id: int, claimed_by: str, worker_status: str, summary_json: str | None = None, error_summary: str | None = None, ) -> dict[str, object]: normalized_claimed_by = normalize_whitespace(claimed_by) if not normalized_claimed_by: raise RetrieverError("claimed_by cannot be empty.") normalized_status = normalize_run_worker_status(worker_status) if normalized_status == "active": raise RetrieverError("finish-run-worker requires a terminal worker status, not 'active'.") summary_payload = decode_json_text(summary_json, default={}) if summary_json is not None else {} if summary_json is not None and not isinstance(summary_payload, dict): raise RetrieverError("summary_json must decode to a JSON object.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN IMMEDIATE") try: run_row = require_run_row_by_id(connection, run_id) worker_row = find_run_worker_row(connection, run_id=run_id, claimed_by=normalized_claimed_by) if worker_row is None: raise RetrieverError( f"Run {run_id} does not have a registered worker for claimed_by={normalized_claimed_by!r}." ) if str(run_row["status"] or "") != "canceled" and normalized_status == "canceled": raise RetrieverError("Worker cannot be finished as canceled unless the run itself is canceled.") already_terminal = ( str(worker_row["status"] or "") == normalized_status and worker_row["completed_at"] is not None ) if already_terminal: payload = { "status": "ok", "idempotent": True, "worker": run_worker_row_to_payload(worker_row), "run": run_status_by_id(connection, run_id), } connection.commit() return payload updated_row = update_run_worker_row( connection, run_id=run_id, claimed_by=normalized_claimed_by, status=normalized_status, last_error=(normalize_whitespace(error_summary) if error_summary is not None else "") or None, summary=summary_payload if isinstance(summary_payload, dict) else {}, completed_at=utc_now(), ) assert updated_row is not None payload = { "status": "ok", "idempotent": False, "worker": run_worker_row_to_payload(updated_row), "run": run_status_by_id(connection, run_id), } connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def finalize_ocr_run( root: Path, *, run_id: int, cascade_metadata: bool = False, metadata_fields: list[str] | None = None, ) -> dict[str, object]: normalized_metadata_fields = normalize_text_derived_metadata_fields(metadata_fields) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN IMMEDIATE") try: materialize_run_items_for_run(connection, paths, root, run_id) payload = finalize_ocr_results_for_run(connection, paths, run_id=run_id) stale_document_ids = sorted( dict.fromkeys(int(document_id) for document_id in payload.get("stale_text_derived_metadata_document_ids") or []) ) payload["metadata_cascade_requested"] = bool(cascade_metadata) payload["next_recommended_commands"] = build_refresh_text_derived_metadata_followup_commands( root, stale_document_ids, field_names=normalized_metadata_fields, ) if cascade_metadata and stale_document_ids: document_rows = fetch_visible_document_rows_by_ids(connection, stale_document_ids) metadata_payload = refresh_text_derived_metadata_for_document_rows( connection, paths, document_rows=document_rows, field_names=normalized_metadata_fields, ) metadata_payload["selector"] = { "mode": "document_ids", "document_ids": [int(row["id"]) for row in document_rows], } metadata_payload["selected_from_scope"] = False metadata_payload["document_ids"] = [int(row["id"]) for row in document_rows] metadata_payload["next_recommended_commands"] = build_rebuild_entities_followup_commands( root, list(metadata_payload.get("entity_rebuild_candidate_document_ids") or []), ) payload["metadata_refresh"] = metadata_payload payload["next_recommended_commands"] = list(metadata_payload.get("next_recommended_commands") or []) connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def finalize_image_description_run(root: Path, *, run_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN") try: materialize_run_items_for_run(connection, paths, root, run_id) payload = finalize_image_description_results_for_run(connection, paths, run_id=run_id) connection.commit() except Exception: connection.rollback() raise return payload finally: connection.close() def publish_run_results( root: Path, *, run_id: int, raw_output_names: list[str] | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) connection.execute("BEGIN") try: publish_summary = publish_result_outputs_for_run( connection, run_id=run_id, output_names=raw_output_names, ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "run": run_summary_by_id(connection, run_id), **publish_summary, } finally: connection.close() def list_jobs(root: Path) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) return { "status": "ok", "jobs": list_job_summaries(connection), } finally: connection.close() def create_job(root: Path, raw_job_name: str, job_kind: str, description: str | None) -> dict[str, object]: job_name = sanitize_processing_identifier(raw_job_name, label="Job name", prefix="job") normalized_kind = normalize_job_kind(job_kind) normalized_description = normalize_whitespace(description) if description and description.strip() else None paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) existing_row = find_job_row_by_name(connection, job_name) if existing_row is not None: raise RetrieverError(f"Job {job_name!r} already exists.") connection.execute("BEGIN") try: job_id = create_job_row( connection, job_name=job_name, job_kind=normalized_kind, description=normalized_description, ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "job": job_summary_by_id(connection, job_id), } finally: connection.close() def add_job_output( root: Path, raw_job_name: str, raw_output_name: str, value_type: str, *, bound_custom_field: str | None = None, description: str | None = None, ) -> dict[str, object]: job_name = sanitize_processing_identifier(raw_job_name, label="Job name", prefix="job") output_name = sanitize_processing_identifier(raw_output_name, label="Job output name", prefix="output") normalized_value_type = normalize_job_output_value_type(value_type) normalized_description = normalize_whitespace(description) if description and description.strip() else None normalized_bound_field = None if bound_custom_field and bound_custom_field.strip(): normalized_bound_field = sanitize_field_name(bound_custom_field) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) job_row = require_job_row_by_name(connection, job_name) connection.execute("BEGIN") try: output_id, created = upsert_job_output_row( connection, job_id=int(job_row["id"]), output_name=output_name, value_type=normalized_value_type, bound_custom_field=normalized_bound_field, description=normalized_description, ) connection.commit() except Exception: connection.rollback() raise output_row = connection.execute( """ SELECT * FROM job_outputs WHERE id = ? """, (output_id,), ).fetchone() assert output_row is not None return { "status": "ok", "created": created, "job": job_summary_by_id(connection, int(job_row["id"])), "job_output": job_output_row_to_payload(output_row), } finally: connection.close() def list_job_versions(root: Path, raw_job_name: str) -> dict[str, object]: job_name = sanitize_processing_identifier(raw_job_name, label="Job name", prefix="job") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) job_row = require_job_row_by_name(connection, job_name) return { "status": "ok", "job": job_summary_by_id(connection, int(job_row["id"])), "job_versions": job_versions_for_job(connection, int(job_row["id"])), } finally: connection.close() def create_job_version( root: Path, raw_job_name: str, *, instruction: str | None, provider: str | None, capability: str | None, model: str | None, input_basis: str | None, response_schema_json: str | None, parameters_json: str | None, segment_profile: str | None, aggregation_strategy: str | None, display_name: str | None, ) -> dict[str, object]: job_name = sanitize_processing_identifier(raw_job_name, label="Job name", prefix="job") normalized_provider = normalize_whitespace(provider) if provider and provider.strip() else "cowork_agent" normalized_instruction = (instruction or "").strip() normalized_model = normalize_whitespace(model) if model and model.strip() else None normalized_segment_profile = ( sanitize_processing_identifier(segment_profile, label="Segment profile", prefix="profile") if segment_profile and segment_profile.strip() else None ) normalized_aggregation = ( sanitize_processing_identifier(aggregation_strategy, label="Aggregation strategy", prefix="aggregation") if aggregation_strategy and aggregation_strategy.strip() else None ) normalized_display_name = normalize_whitespace(display_name) if display_name and display_name.strip() else None parsed_response_schema = parse_json_argument( response_schema_json, label="Response schema", default=None, ) response_schema_text = None if parsed_response_schema is None else compact_json_text(parsed_response_schema) parameters = parse_json_object_argument( parameters_json, label="Parameters", default={}, ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) job_row = require_job_row_by_name(connection, job_name) job_kind = normalize_job_kind(str(job_row["job_kind"])) normalized_capability = ( normalize_job_capability(capability) if capability and capability.strip() else default_job_capability_for_kind(job_kind) ) normalized_input_basis = ( normalize_job_input_basis(input_basis) if input_basis and input_basis.strip() else default_job_input_basis_for_kind(job_kind) ) if job_kind == "translation" or normalized_capability == "text_translation": target_language = normalize_whitespace( str(parameters.get("target_language") or parameters.get("target_lang") or parameters.get("language") or "") ).lower() if not target_language: raise RetrieverError("Translation job versions require parameters_json.target_language.") connection.execute("BEGIN") try: version_id = create_job_version_row( connection, job_id=int(job_row["id"]), job_name=job_name, instruction_text=normalized_instruction, response_schema_json=response_schema_text, capability=normalized_capability, provider=normalized_provider, model=normalized_model, parameters_json=compact_json_text(parameters), input_basis=normalized_input_basis, segment_profile=normalized_segment_profile, aggregation_strategy=normalized_aggregation, display_name=normalized_display_name, ) connection.commit() except Exception: connection.rollback() raise version_row = connection.execute( """ SELECT * FROM job_versions WHERE id = ? """, (version_id,), ).fetchone() assert version_row is not None return { "status": "ok", "job": job_summary_by_id(connection, int(job_row["id"])), "job_version": job_version_row_to_payload(version_row), } finally: connection.close() def get_custom_field_registry_row(connection: sqlite3.Connection, field_name: str) -> sqlite3.Row | None: return connection.execute( """ SELECT field_name, field_type, instruction, created_at FROM custom_fields_registry WHERE field_name = ? """, (field_name,), ).fetchone() def require_custom_field_registry_row(connection: sqlite3.Connection, field_name: str) -> sqlite3.Row: registry_row = get_custom_field_registry_row(connection, field_name) if registry_row is None: raise RetrieverError(f"Unknown custom field: {field_name}") return registry_row def count_documents_with_non_null_field_value(connection: sqlite3.Connection, field_name: str) -> int: row = connection.execute( f""" SELECT COUNT(*) AS value_count FROM documents WHERE {quote_identifier(field_name)} IS NOT NULL """, ).fetchone() return int(row["value_count"]) if row is not None else 0 def list_fields(root: Path) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) columns = table_columns(connection, "documents") fields: list[dict[str, object]] = [] rows = connection.execute( """ SELECT field_name, field_type, instruction, created_at FROM custom_fields_registry ORDER BY field_name ASC """ ).fetchall() for row in rows: field_name = str(row["field_name"]) if field_name not in columns: continue fields.append( { "field_name": field_name, "field_type": str(row["field_type"]), "instruction": row["instruction"], "created_at": row["created_at"], "documents_with_values": count_documents_with_non_null_field_value(connection, field_name), } ) return {"status": "ok", "fields": fields} finally: connection.close() def describe_field( root: Path, raw_field_name: str, *, text: str | None = None, clear: bool = False, ) -> dict[str, object]: if clear and text is not None: raise RetrieverError("Choose either --text or --clear, not both.") if not clear and text is None: raise RetrieverError("Provide --text or --clear.") normalized_field_name = sanitize_field_name(raw_field_name) next_instruction = None if clear else text paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) require_custom_field_registry_row(connection, normalized_field_name) connection.execute( """ UPDATE custom_fields_registry SET instruction = ? WHERE field_name = ? """, (next_instruction, normalized_field_name), ) connection.commit() updated_row = require_custom_field_registry_row(connection, normalized_field_name) return { "status": "ok", "field_name": normalized_field_name, "field_type": str(updated_row["field_type"]), "instruction": updated_row["instruction"], } finally: connection.close() def ensure_fill_target_field_definition(connection: sqlite3.Connection, field_name: str) -> dict[str, str]: try: field_def = resolve_field_definition(connection, field_name) except RetrieverError as exc: suggestions = field_name_suggestions(connection, field_name) suggestion_text = f" Did you mean: {', '.join(suggestions)}?" if suggestions else "" raise RetrieverError(f"Unknown field '{field_name}'.{suggestion_text}") from exc canonical_name = str(field_def["field_name"]) source = str(field_def.get("source") or "") if source == "virtual": raise RetrieverError(f"Field '{canonical_name}' is derived and cannot be filled manually.") if canonical_name in SYSTEM_MANAGED_FIELDS: raise RetrieverError(f"Field '{canonical_name}' is system-managed and cannot be filled manually.") if source == "builtin": if canonical_name not in EDITABLE_BUILTIN_FIELDS: raise RetrieverError(f"Field '{canonical_name}' is not an editable built-in field.") return field_def if source == "custom": return field_def raise RetrieverError( f"Field '{canonical_name}' is not a registered custom field or editable built-in field." ) def replace_document_field_locks( connection: sqlite3.Connection, old_field_name: str, new_field_name: str, ) -> int: rows = connection.execute( f""" SELECT id, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} AS locks_json FROM documents WHERE {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} LIKE ? """, (f'%"{old_field_name}"%',), ).fetchall() updated = 0 for row in rows: locks = normalize_string_list(row["locks_json"]) if old_field_name not in locks: continue next_locks: list[str] = [] for lock_name in locks: target_name = new_field_name if lock_name == old_field_name else lock_name if target_name not in next_locks: next_locks.append(target_name) connection.execute( f""" UPDATE documents SET {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} = ? WHERE id = ? """, (json.dumps(next_locks), int(row["id"])), ) updated += 1 return updated def sample_documents_for_fill( connection: sqlite3.Connection, document_ids: list[int], field_name: str, *, limit: int = 5, ) -> list[dict[str, object]]: normalized_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids))[:limit] if not normalized_ids: return [] placeholders = ", ".join("?" for _ in normalized_ids) rows = connection.execute( f""" SELECT id, control_number, file_name, title, {quote_identifier(field_name)} AS field_value FROM documents WHERE id IN ({placeholders}) """, tuple(normalized_ids), ).fetchall() rows_by_id = {int(row["id"]): row for row in rows} sample_rows: list[dict[str, object]] = [] for document_id in normalized_ids: row = rows_by_id.get(document_id) if row is None: continue sample_rows.append( { "document_id": int(row["id"]), "control_number": row["control_number"], "file_name": row["file_name"], "title": row["title"], "value": row["field_value"], } ) return sample_rows def normalize_delete_path_prefixes(raw_path_prefixes: list[str] | None) -> list[str]: normalized_prefixes: list[str] = [] for raw_path_prefix in raw_path_prefixes or []: candidate = str(raw_path_prefix or "").strip() if candidate.startswith("./"): candidate = candidate[2:] candidate = candidate.rstrip("/") if candidate: normalized_prefixes.append(candidate) return list(dict.fromkeys(normalized_prefixes)) def sql_text_literal(value: str) -> str: return "'" + value.replace("'", "''") + "'" def path_prefix_scope_expression(path_prefixes: list[str]) -> str | None: clauses: list[str] = [] for path_prefix in normalize_delete_path_prefixes(path_prefixes): clauses.append( f"(rel_path = {sql_text_literal(path_prefix)} OR rel_path LIKE {sql_text_literal(path_prefix + '/%')})" ) if not clauses: return None return " OR ".join(clauses) def raw_filters_include_occurrence_scope( connection: sqlite3.Connection, raw_filters: object | None, ) -> bool: if not raw_filters: return False pattern = re.compile( r"\b(?:" + "|".join(re.escape(field_name) for field_name in sorted(OCCURRENCE_FILTER_FIELDS, key=len, reverse=True)) + r")\b", re.IGNORECASE, ) return any( pattern.search(expression) is not None for expression in normalize_sql_filter_expressions(raw_filters) ) def raw_filter_item_includes_occurrence_scope( connection: sqlite3.Connection, raw_filter_item: object, ) -> bool: if isinstance(raw_filter_item, str): expression = raw_filter_item elif isinstance(raw_filter_item, (list, tuple)): expression = " ".join( str(part) for part in raw_filter_item if normalize_inline_whitespace(str(part or "")) ) else: return False if not expression: return False pattern = re.compile( r"\b(?:" + "|".join(re.escape(field_name) for field_name in sorted(OCCURRENCE_FILTER_FIELDS, key=len, reverse=True)) + r")\b", re.IGNORECASE, ) return pattern.search(expression) is not None def document_only_raw_filters( connection: sqlite3.Connection, raw_filters: object | None, ) -> list[object]: document_filters: list[object] = [] for raw_filter_item in raw_filters or []: if raw_filter_item_includes_occurrence_scope(connection, raw_filter_item): continue document_filters.append(raw_filter_item) return document_filters def document_ids_matching_occurrence_filters( connection: sqlite3.Connection, raw_filters: object | None, ) -> list[int]: occurrence_scope_clauses, occurrence_scope_params = build_occurrence_scope_filters(connection, raw_filters) rows = connection.execute( f""" SELECT DISTINCT d.id FROM document_occurrences o JOIN documents d ON d.id = o.document_id WHERE d.lifecycle_status NOT IN ('missing', 'deleted') AND {' AND '.join(occurrence_scope_clauses)} ORDER BY d.id ASC """, occurrence_scope_params, ).fetchall() return [int(row["id"]) for row in rows] def fetch_deletable_document_rows_by_ids( connection: sqlite3.Connection, document_ids: list[int], ) -> list[sqlite3.Row]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids)) if not normalized_document_ids: return [] placeholders = ", ".join("?" for _ in normalized_document_ids) rows = connection.execute( f""" SELECT * FROM documents WHERE id IN ({placeholders}) """, normalized_document_ids, ).fetchall() rows_by_id = {int(row["id"]): row for row in rows} missing_ids: list[int] = [] lifecycle_hidden: list[str] = [] visible_rows_by_id: dict[int, sqlite3.Row] = {} for document_id in normalized_document_ids: row = rows_by_id.get(document_id) if row is None: missing_ids.append(document_id) continue if row["lifecycle_status"] in {"missing", "deleted"}: lifecycle_hidden.append(f"{document_id} ({row['lifecycle_status']})") continue visible_rows_by_id[document_id] = row errors: list[str] = [] if missing_ids: errors.append( "Unknown document id" + ("" if len(missing_ids) == 1 else "s") + f": {', '.join(str(document_id) for document_id in missing_ids)}" ) if lifecycle_hidden: errors.append( "Document id" + ("" if len(lifecycle_hidden) == 1 else "s") + " not deletable due to lifecycle_status: " + ", ".join(lifecycle_hidden) ) if errors: raise RetrieverError(" ".join(errors)) return [visible_rows_by_id[document_id] for document_id in normalized_document_ids] def sample_documents_for_delete( connection: sqlite3.Connection, document_ids: list[int], *, limit: int = 5, ) -> list[dict[str, object]]: normalized_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids))[:limit] if not normalized_ids: return [] placeholders = ", ".join("?" for _ in normalized_ids) rows = connection.execute( f""" SELECT id, control_number, rel_path, file_name, title FROM documents WHERE id IN ({placeholders}) """, tuple(normalized_ids), ).fetchall() rows_by_id = {int(row["id"]): row for row in rows} sample_rows: list[dict[str, object]] = [] for document_id in normalized_ids: row = rows_by_id.get(document_id) if row is None: continue sample_rows.append( { "document_id": int(row["id"]), "control_number": row["control_number"], "rel_path": row["rel_path"], "file_name": row["file_name"], "title": row["title"], } ) return sample_rows def scoped_active_occurrence_rows_for_document( connection: sqlite3.Connection, document_id: int, occurrence_scope_clauses: list[str], occurrence_scope_params: list[object], ) -> list[sqlite3.Row]: return connection.execute( f""" SELECT o.* FROM document_occurrences o JOIN documents d ON d.id = o.document_id WHERE o.document_id = ? AND {' AND '.join(occurrence_scope_clauses)} ORDER BY o.id ASC """, [document_id, *occurrence_scope_params], ).fetchall() def collect_occurrence_rows_for_delete_plan( connection: sqlite3.Connection, *, root_occurrence_ids: set[int], direct_occurrence_ids: set[int], ) -> list[sqlite3.Row]: rows_by_id: dict[int, sqlite3.Row] = {} normalized_root_ids = sorted(int(occurrence_id) for occurrence_id in root_occurrence_ids) if normalized_root_ids: placeholders = ", ".join("?" for _ in normalized_root_ids) root_rows = connection.execute( f""" SELECT * FROM document_occurrences WHERE id IN ({placeholders}) OR parent_occurrence_id IN ({placeholders}) ORDER BY id ASC """, [*normalized_root_ids, *normalized_root_ids], ).fetchall() for row in root_rows: rows_by_id[int(row["id"])] = row normalized_direct_ids = sorted(int(occurrence_id) for occurrence_id in direct_occurrence_ids) if normalized_direct_ids: placeholders = ", ".join("?" for _ in normalized_direct_ids) direct_rows = connection.execute( f""" SELECT * FROM document_occurrences WHERE id IN ({placeholders}) ORDER BY id ASC """, normalized_direct_ids, ).fetchall() for row in direct_rows: rows_by_id[int(row["id"])] = row return [rows_by_id[occurrence_id] for occurrence_id in sorted(rows_by_id)] def delete_documents_with_no_active_occurrences( connection: sqlite3.Connection, paths: dict[str, Path], document_ids: list[int] | set[int], *, deleted_at: str, ) -> list[int]: deleted_document_ids: list[int] = [] for document_id in sorted({int(value) for value in document_ids}): remaining_row = connection.execute( """ SELECT 1 FROM document_occurrences WHERE document_id = ? AND lifecycle_status = 'active' LIMIT 1 """, (document_id,), ).fetchone() if remaining_row is not None: continue document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (document_id,), ).fetchone() if document_row is None or document_row["lifecycle_status"] == "deleted": continue cleanup_document_artifacts(paths, connection, document_row) delete_document_related_rows(connection, document_id) connection.execute( """ UPDATE documents SET dataset_id = NULL, lifecycle_status = 'deleted', updated_at = ? WHERE id = ? """, (deleted_at, document_id), ) deleted_document_ids.append(document_id) return deleted_document_ids def delete_dataset_memberships_for_documents( connection: sqlite3.Connection, document_ids: list[int] | set[int], ) -> int: normalized_document_ids = sorted({int(document_id) for document_id in document_ids}) if not normalized_document_ids: return 0 placeholders = ", ".join("?" for _ in normalized_document_ids) cursor = connection.execute( f""" DELETE FROM dataset_documents WHERE document_id IN ({placeholders}) """, normalized_document_ids, ) return int(cursor.rowcount or 0) def round_half_away_from_zero(raw_value: float) -> int: if raw_value >= 0: return int(raw_value + 0.5) return -int(abs(raw_value) + 0.5) def convert_field_value_for_type_change( raw_value: object, from_type: str, to_type: str, ) -> tuple[object, bool]: if raw_value is None: return None, False normalized_from_type = from_type.strip().lower() normalized_to_type = to_type.strip().lower() if normalized_from_type == normalized_to_type: return raw_value, False if normalized_from_type == "text": raw_text = str(raw_value) if not raw_text.strip(): return None, True if normalized_to_type == "text": return raw_text, False converted_value = value_from_type(normalized_to_type, raw_text) return converted_value, raw_text != str(converted_value) if normalized_from_type == "date": if normalized_to_type != "text": raise RetrieverError(f"Cannot convert {normalized_from_type} to {normalized_to_type}.") return str(raw_value), False if normalized_from_type == "boolean": int_value = int(raw_value) if normalized_to_type == "integer": return int_value, False if normalized_to_type == "real": return float(int_value), True if normalized_to_type == "text": return ("true" if int_value else "false"), True raise RetrieverError(f"Cannot convert {normalized_from_type} to {normalized_to_type}.") if normalized_from_type == "integer": int_value = int(raw_value) if normalized_to_type == "boolean": if int_value not in {0, 1}: raise RetrieverError(f"Expected 0 or 1 for boolean conversion, got {int_value!r}") return int_value, False if normalized_to_type == "real": return float(int_value), True if normalized_to_type == "text": return str(int_value), True raise RetrieverError(f"Cannot convert {normalized_from_type} to {normalized_to_type}.") if normalized_from_type == "real": float_value = float(raw_value) if normalized_to_type == "integer": rounded_value = round_half_away_from_zero(float_value) return rounded_value, rounded_value != float_value if normalized_to_type == "text": return str(float_value), True raise RetrieverError(f"Cannot convert {normalized_from_type} to {normalized_to_type}.") raise RetrieverError(f"Unsupported field type conversion: {normalized_from_type} -> {normalized_to_type}") def field_type_conversion_allowed(from_type: str, to_type: str) -> bool: normalized_from = from_type.strip().lower() normalized_to = to_type.strip().lower() if normalized_from == normalized_to: return True allowed_pairs = { ("boolean", "integer"), ("boolean", "real"), ("boolean", "text"), ("date", "text"), ("integer", "boolean"), ("integer", "real"), ("integer", "text"), ("real", "integer"), ("real", "text"), ("text", "boolean"), ("text", "date"), ("text", "integer"), ("text", "real"), } return (normalized_from, normalized_to) in allowed_pairs def build_no_document_selection_error() -> str: return ( "No document selection active. Provide 'on ', narrow the current browse state " "with /dataset / /filter / /search / /bates / /from-run, or use fill-field with explicit selectors." ) def add_field(root: Path, raw_field_name: str, field_type: str, instruction: str | None) -> dict[str, object]: normalized_field_name = sanitize_field_name(raw_field_name) normalized_field_type = field_type.strip().lower() if normalized_field_type not in REGISTRY_FIELD_TYPES: raise RetrieverError(f"Unsupported field type: {field_type}") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) columns = table_columns(connection, "documents") existing_registry_row = get_custom_field_registry_row(connection, normalized_field_name) if existing_registry_row is not None and existing_registry_row["field_type"] != normalized_field_type: if existing_registry_row["field_type"] == "text" and normalized_field_type == "date": raise RetrieverError( f"Field '{normalized_field_name}' already exists as text; use promote-field-type to validate and promote it to date." ) raise RetrieverError( f"Field '{normalized_field_name}' already exists with type {existing_registry_row['field_type']!r}." ) sql_type = REGISTRY_FIELD_TYPES[normalized_field_type] if normalized_field_name not in columns: connection.execute( f"ALTER TABLE documents ADD COLUMN {quote_identifier(normalized_field_name)} {sql_type}" ) connection.execute( """ INSERT INTO custom_fields_registry (field_name, field_type, instruction, created_at) VALUES (?, ?, ?, ?) ON CONFLICT(field_name) DO UPDATE SET field_type = excluded.field_type, instruction = excluded.instruction """, (normalized_field_name, normalized_field_type, instruction, utc_now()), ) connection.commit() registry_status = reconcile_custom_fields_registry(connection, repair=True) return { "status": "ok", "field_name": normalized_field_name, "field_type": normalized_field_type, "instruction": instruction, "custom_field_registry": registry_status, } finally: connection.close() def rename_field(root: Path, old_name: str, new_name: str) -> dict[str, object]: normalized_old_name = sanitize_field_name(old_name) normalized_new_name = sanitize_field_name(new_name) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) require_custom_field_registry_row(connection, normalized_old_name) if normalized_old_name not in table_columns(connection, "documents"): raise RetrieverError(f"Field column '{normalized_old_name}' does not exist on documents.") if normalized_old_name == normalized_new_name: return { "status": "ok", "renamed_from": normalized_old_name, "field_name": normalized_new_name, "locks_updated": 0, "state_updates": {}, } if normalized_new_name in table_columns(connection, "documents") or get_custom_field_registry_row( connection, normalized_new_name, ) is not None: raise RetrieverError(f"Field '{normalized_new_name}' already exists.") state_plan = plan_field_rename_state_changes(paths, normalized_old_name, normalized_new_name) blockers = state_plan.get("blockers") if isinstance(blockers, list) and blockers: return { "status": "blocked", "renamed_from": normalized_old_name, "field_name": normalized_new_name, "blockers": blockers, } connection.execute("BEGIN") try: connection.execute( f""" ALTER TABLE documents RENAME COLUMN {quote_identifier(normalized_old_name)} TO {quote_identifier(normalized_new_name)} """ ) connection.execute( """ UPDATE custom_fields_registry SET field_name = ? WHERE field_name = ? """, (normalized_new_name, normalized_old_name), ) locks_updated = replace_document_field_locks(connection, normalized_old_name, normalized_new_name) connection.commit() except Exception: connection.rollback() raise apply_field_state_change_plan(paths, state_plan) return { "status": "ok", "renamed_from": normalized_old_name, "field_name": normalized_new_name, "locks_updated": locks_updated, "state_updates": state_plan.get("changes") or {}, } finally: connection.close() def delete_field(root: Path, raw_field_name: str, *, confirm: bool = False) -> dict[str, object]: normalized_field_name = sanitize_field_name(raw_field_name) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) require_custom_field_registry_row(connection, normalized_field_name) if normalized_field_name not in table_columns(connection, "documents"): raise RetrieverError(f"Field column '{normalized_field_name}' does not exist on documents.") non_null_values_removed = count_documents_with_non_null_field_value(connection, normalized_field_name) preview_payload = { "field_name": normalized_field_name, "non_null_values_removed": non_null_values_removed, "documents_affected": non_null_values_removed, } state_plan = plan_field_delete_state_changes(paths, normalized_field_name) blockers = state_plan.get("blockers") if isinstance(blockers, list) and blockers: result = {"status": "blocked", **preview_payload, "blockers": blockers} pending_changes = state_plan.get("changes") if pending_changes: result["state_updates"] = pending_changes return result if not confirm: return { "status": "confirm_required", **preview_payload, "state_updates": state_plan.get("changes") or {}, } connection.execute("BEGIN") try: locks_removed = drop_document_field_locks(connection, normalized_field_name) connection.execute( f"ALTER TABLE documents DROP COLUMN {quote_identifier(normalized_field_name)}" ) connection.execute( """ DELETE FROM custom_fields_registry WHERE field_name = ? """, (normalized_field_name,), ) connection.commit() except Exception: connection.rollback() raise apply_field_state_change_plan(paths, state_plan) return { "status": "ok", "deleted": normalized_field_name, "non_null_values_removed": non_null_values_removed, "documents_affected": non_null_values_removed, "locks_removed": locks_removed, "state_updates": state_plan.get("changes") or {}, } finally: connection.close() def change_field_type(root: Path, raw_field_name: str, target_field_type: str) -> dict[str, object]: normalized_field_name = sanitize_field_name(raw_field_name) normalized_target_type = target_field_type.strip().lower() if normalized_target_type not in REGISTRY_FIELD_TYPES: raise RetrieverError(f"Unsupported field type: {target_field_type}") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) registry_row = require_custom_field_registry_row(connection, normalized_field_name) current_type = str(registry_row["field_type"]) if current_type == normalized_target_type: return { "status": "ok", "field_name": normalized_field_name, "from_type": current_type, "to_type": normalized_target_type, "normalized_values_updated": 0, "documents_checked": 0, "documents_with_values": 0, "conversion_applied": False, "promotion_applied": False, } if not field_type_conversion_allowed(current_type, normalized_target_type): raise RetrieverError(f"Field '{normalized_field_name}' cannot convert from {current_type!r} to {normalized_target_type!r}.") if normalized_field_name not in table_columns(connection, "documents"): raise RetrieverError(f"Field column '{normalized_field_name}' does not exist on documents.") value_rows = connection.execute( f""" SELECT id, {quote_identifier(normalized_field_name)} AS value FROM documents WHERE {quote_identifier(normalized_field_name)} IS NOT NULL ORDER BY id ASC """ ).fetchall() invalid_values: list[dict[str, object]] = [] normalized_updates: list[tuple[object, int]] = [] warnings: list[str] = ( ["real -> integer rounds values to the nearest integer (half away from zero)."] if current_type == "real" and normalized_target_type == "integer" else [] ) for row in value_rows: raw_value = row["value"] try: normalized_value, changed = convert_field_value_for_type_change( raw_value, current_type, normalized_target_type, ) except RetrieverError: if len(invalid_values) < 10: invalid_values.append({"document_id": int(row["id"]), "value": raw_value}) continue if changed or normalized_value != raw_value: normalized_updates.append((normalized_value, int(row["id"]))) if invalid_values: return { "status": "blocked", "field_name": normalized_field_name, "from_type": current_type, "to_type": normalized_target_type, "documents_checked": len(value_rows), "documents_with_values": len(value_rows), "invalid_value_samples": invalid_values, "conversion_applied": False, "promotion_applied": False, } source_sql_type = REGISTRY_FIELD_TYPES[current_type] target_sql_type = REGISTRY_FIELD_TYPES[normalized_target_type] conversion_requires_column_rebuild = source_sql_type != target_sql_type temp_column_name = f"{normalized_field_name}__tmp_{normalized_target_type}" connection.execute("BEGIN") try: if conversion_requires_column_rebuild: if temp_column_name in table_columns(connection, "documents"): raise RetrieverError(f"Temporary field column '{temp_column_name}' already exists.") connection.execute( f"ALTER TABLE documents ADD COLUMN {quote_identifier(temp_column_name)} {target_sql_type}" ) if value_rows: for row in value_rows: document_id = int(row["id"]) raw_value = row["value"] converted_value, _ = convert_field_value_for_type_change( raw_value, current_type, normalized_target_type, ) connection.execute( f""" UPDATE documents SET {quote_identifier(temp_column_name)} = ?, updated_at = ? WHERE id = ? """, (converted_value, utc_now(), document_id), ) connection.execute( f"ALTER TABLE documents DROP COLUMN {quote_identifier(normalized_field_name)}" ) connection.execute( f""" ALTER TABLE documents RENAME COLUMN {quote_identifier(temp_column_name)} TO {quote_identifier(normalized_field_name)} """ ) else: for normalized_value, document_id in normalized_updates: connection.execute( f""" UPDATE documents SET {quote_identifier(normalized_field_name)} = ?, updated_at = ? WHERE id = ? """, (normalized_value, utc_now(), document_id), ) connection.execute( """ UPDATE custom_fields_registry SET field_type = ? WHERE field_name = ? """, (normalized_target_type, normalized_field_name), ) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "field_name": normalized_field_name, "from_type": current_type, "to_type": normalized_target_type, "documents_checked": len(value_rows), "documents_with_values": len(value_rows), "normalized_values_updated": len(normalized_updates), "conversion_applied": True, "promotion_applied": True, "column_rebuilt": conversion_requires_column_rebuild, "warnings": warnings, } finally: connection.close() def promote_field_type(root: Path, raw_field_name: str, target_field_type: str) -> dict[str, object]: if target_field_type.strip().lower() != "date": raise RetrieverError("Only text -> date field promotion is supported via promote-field-type.") payload = change_field_type(root, raw_field_name, target_field_type) payload.setdefault("promotion_applied", bool(payload.get("conversion_applied"))) return payload def get_document_row_for_conversation_assignment(connection: sqlite3.Connection, document_id: int) -> sqlite3.Row: row = connection.execute( """ SELECT id, parent_document_id, rel_path, control_number, content_type, source_kind, source_rel_path, source_folder_path, title, subject, conversation_id, conversation_assignment_mode, lifecycle_status FROM documents WHERE id = ? """, (document_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown document id: {document_id}") if row["lifecycle_status"] in {"missing", "deleted"}: raise RetrieverError(f"Document {document_id} is not active.") return row def get_document_family_root_row_for_assignment(connection: sqlite3.Connection, document_id: int) -> sqlite3.Row: row = get_document_row_for_conversation_assignment(connection, document_id) while row["parent_document_id"] is not None: row = get_document_row_for_conversation_assignment(connection, int(row["parent_document_id"])) return row def document_conversation_assignment_category(row: sqlite3.Row) -> str | None: content_type = normalize_whitespace(str(row["content_type"] or "")) source_kind = normalize_whitespace(str(row["source_kind"] or "")).lower() if content_type == "Email": return "email" if content_type == "Chat" and source_kind == PST_SOURCE_KIND: return "pst_chat" return None def ensure_document_supports_manual_conversation_assignment(row: sqlite3.Row) -> str: category = document_conversation_assignment_category(row) if category is None: raise RetrieverError( "Manual conversation changes currently support top-level email documents and PST chat documents only." ) return category def manual_conversation_display_name( root_row: sqlite3.Row, *, category: str, existing_conversation_row: sqlite3.Row | None, ) -> str: if category == "email": for candidate in (root_row["subject"], root_row["title"]): display_name = normalize_email_thread_subject(candidate, preserve_case=True) if display_name: return display_name if existing_conversation_row is not None: existing_display = normalize_whitespace(str(existing_conversation_row["display_name"] or "")) if existing_display: return existing_display return "Email conversation" title = normalize_whitespace(str(root_row["title"] or "")) if title: return title folder_path = normalize_whitespace(str(root_row["source_folder_path"] or "")) if folder_path: leaf_name = normalize_whitespace(folder_path.split("/")[-1]) if leaf_name: return leaf_name if existing_conversation_row is not None: existing_display = normalize_whitespace(str(existing_conversation_row["display_name"] or "")) if existing_display: return existing_display return "Chat conversation" def create_manual_singleton_conversation( connection: sqlite3.Connection, root_row: sqlite3.Row, ) -> int: category = ensure_document_supports_manual_conversation_assignment(root_row) existing_conversation_row = None if root_row["conversation_id"] is not None: existing_conversation_row = connection.execute( "SELECT * FROM conversations WHERE id = ?", (int(root_row["conversation_id"]),), ).fetchone() if existing_conversation_row is not None: source_kind = str(existing_conversation_row["source_kind"]) source_locator = str(existing_conversation_row["source_locator"]) conversation_type = str(existing_conversation_row["conversation_type"]) elif category == "email": source_kind = EMAIL_CONVERSATION_SOURCE_KIND source_locator = filesystem_dataset_locator() conversation_type = "email" else: source_kind = PST_SOURCE_KIND source_locator = normalize_whitespace(str(root_row["source_rel_path"] or "")) or filesystem_dataset_locator() conversation_type = "chat" return upsert_conversation_row( connection, source_kind=source_kind, source_locator=source_locator, conversation_key=f"manual:{category}:{int(root_row['id'])}", conversation_type=conversation_type, display_name=manual_conversation_display_name( root_row, category=category, existing_conversation_row=existing_conversation_row, ), ) def reassign_conversations_and_refresh_previews( connection: sqlite3.Connection, paths: dict[str, Path], ) -> dict[str, int]: assignment = assign_supported_conversations(connection) refresh_conversation_previews(connection, paths) return assignment def resolve_conversation_preview_refresh_ids( connection: sqlite3.Connection, *, conversation_ids: list[int] | None = None, document_ids: list[int] | None = None, dataset_id: int | None = None, dataset_name: str | None = None, ignore_documents_without_conversation: bool = False, ) -> tuple[list[int], dict[str, object] | None]: target_conversation_ids: set[int] = set() dataset_summary: dict[str, object] | None = None if conversation_ids: requested_conversation_ids = sorted(dict.fromkeys(int(conversation_id) for conversation_id in conversation_ids)) rows = connection.execute( f""" SELECT id FROM conversations WHERE id IN ({", ".join("?" for _ in requested_conversation_ids)}) """, tuple(requested_conversation_ids), ).fetchall() found_conversation_ids = {int(row["id"]) for row in rows} missing_conversation_ids = [ conversation_id for conversation_id in requested_conversation_ids if conversation_id not in found_conversation_ids ] if missing_conversation_ids: missing_text = ", ".join(str(conversation_id) for conversation_id in missing_conversation_ids) raise RetrieverError(f"Unknown conversation id(s): {missing_text}") target_conversation_ids.update(requested_conversation_ids) if document_ids: for document_id in sorted(dict.fromkeys(int(document_id) for document_id in document_ids)): root_row = get_document_family_root_row_for_assignment(connection, document_id) conversation_id = root_row["conversation_id"] if conversation_id is None: if ignore_documents_without_conversation: continue raise RetrieverError( f"Document {document_id} does not belong to a conversation, so there are no conversation previews to refresh." ) target_conversation_ids.add(int(conversation_id)) if dataset_id is not None or dataset_name is not None: dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) dataset_summary = dataset_summary_by_id(connection, int(dataset_row["id"])) rows = connection.execute( """ SELECT DISTINCT documents.conversation_id AS conversation_id FROM dataset_documents JOIN documents ON documents.id = dataset_documents.document_id WHERE dataset_documents.dataset_id = ? AND documents.conversation_id IS NOT NULL AND documents.lifecycle_status NOT IN ('missing', 'deleted') ORDER BY documents.conversation_id ASC """, (int(dataset_row["id"]),), ).fetchall() target_conversation_ids.update( int(row["conversation_id"]) for row in rows if row["conversation_id"] is not None ) if not target_conversation_ids and conversation_ids is None and document_ids is None and dataset_id is None and dataset_name is None: target_conversation_ids.update(list_active_conversation_ids(connection)) return sorted(target_conversation_ids), dataset_summary def filter_conversation_ids_with_missing_preview_artifacts( connection: sqlite3.Connection, paths: dict[str, Path], conversation_ids: list[int], ) -> list[int]: normalized_conversation_ids = sorted(dict.fromkeys(int(conversation_id) for conversation_id in conversation_ids)) if not normalized_conversation_ids: return [] rows = connection.execute( f""" SELECT d.conversation_id, d.id AS document_id, dp.rel_preview_path, dp.preview_type FROM documents d LEFT JOIN document_previews dp ON dp.document_id = d.id WHERE d.conversation_id IN ({", ".join("?" for _ in normalized_conversation_ids)}) AND d.lifecycle_status NOT IN ('missing', 'deleted') AND COALESCE(d.child_document_kind, '') != ? ORDER BY d.conversation_id ASC, d.id ASC, dp.ordinal ASC, dp.id ASC """, (*normalized_conversation_ids, CHILD_DOCUMENT_KIND_ATTACHMENT), ).fetchall() missing_conversation_ids: set[int] = set() for row in rows: conversation_id = int(row["conversation_id"]) rel_preview_path = row["rel_preview_path"] if rel_preview_path is None: missing_conversation_ids.add(conversation_id) continue if normalize_whitespace(str(row["preview_type"] or "")).lower() == "native": continue if not (paths["state_dir"] / str(rel_preview_path)).exists(): missing_conversation_ids.add(conversation_id) return [ conversation_id for conversation_id in normalized_conversation_ids if conversation_id in missing_conversation_ids ] def resolve_document_preview_refresh_ids( connection: sqlite3.Connection, *, conversation_ids: list[int] | None = None, document_ids: list[int] | None = None, dataset_id: int | None = None, dataset_name: str | None = None, ) -> tuple[list[int], dict[str, object] | None]: target_document_ids: set[int] = set() dataset_summary: dict[str, object] | None = None if conversation_ids: requested_conversation_ids = sorted(dict.fromkeys(int(conversation_id) for conversation_id in conversation_ids)) rows = connection.execute( f""" SELECT id FROM conversations WHERE id IN ({", ".join("?" for _ in requested_conversation_ids)}) """, tuple(requested_conversation_ids), ).fetchall() found_conversation_ids = {int(row["id"]) for row in rows} missing_conversation_ids = [ conversation_id for conversation_id in requested_conversation_ids if conversation_id not in found_conversation_ids ] if missing_conversation_ids: missing_text = ", ".join(str(conversation_id) for conversation_id in missing_conversation_ids) raise RetrieverError(f"Unknown conversation id(s): {missing_text}") rows = connection.execute( f""" SELECT id FROM documents WHERE conversation_id IN ({", ".join("?" for _ in requested_conversation_ids)}) AND lifecycle_status NOT IN ('missing', 'deleted') ORDER BY id ASC """, tuple(requested_conversation_ids), ).fetchall() target_document_ids.update(int(row["id"]) for row in rows) if document_ids: requested_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) rows = connection.execute( f""" SELECT id FROM documents WHERE id IN ({", ".join("?" for _ in requested_document_ids)}) AND lifecycle_status NOT IN ('missing', 'deleted') """, tuple(requested_document_ids), ).fetchall() found_document_ids = {int(row["id"]) for row in rows} missing_document_ids = [ document_id for document_id in requested_document_ids if document_id not in found_document_ids ] if missing_document_ids: missing_text = ", ".join(str(document_id) for document_id in missing_document_ids) raise RetrieverError(f"Unknown active document id(s): {missing_text}") target_document_ids.update(requested_document_ids) if dataset_id is not None or dataset_name is not None: dataset_row = resolve_dataset_row(connection, dataset_id=dataset_id, dataset_name=dataset_name) dataset_summary = dataset_summary_by_id(connection, int(dataset_row["id"])) rows = connection.execute( """ SELECT DISTINCT documents.id AS document_id FROM dataset_documents JOIN documents ON documents.id = dataset_documents.document_id WHERE dataset_documents.dataset_id = ? AND documents.lifecycle_status NOT IN ('missing', 'deleted') ORDER BY documents.id ASC """, (int(dataset_row["id"]),), ).fetchall() target_document_ids.update(int(row["document_id"]) for row in rows) if ( not target_document_ids and conversation_ids is None and document_ids is None and dataset_id is None and dataset_name is None ): rows = connection.execute( """ SELECT id FROM documents WHERE lifecycle_status NOT IN ('missing', 'deleted') ORDER BY id ASC """ ).fetchall() target_document_ids.update(int(row["id"]) for row in rows) return sorted(target_document_ids), dataset_summary def filter_document_ids_with_missing_preview_artifacts( connection: sqlite3.Connection, paths: dict[str, Path], document_ids: list[int], ) -> list[int]: normalized_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) if not normalized_document_ids: return [] rows = connection.execute( f""" SELECT d.id AS document_id, dp.rel_preview_path, dp.preview_type FROM documents d LEFT JOIN document_previews dp ON dp.document_id = d.id WHERE d.id IN ({", ".join("?" for _ in normalized_document_ids)}) AND d.lifecycle_status NOT IN ('missing', 'deleted') ORDER BY d.id ASC, dp.ordinal ASC, dp.id ASC """, tuple(normalized_document_ids), ).fetchall() missing_document_ids: set[int] = set() for row in rows: if row["rel_preview_path"] is None: missing_document_ids.add(int(row["document_id"])) continue if normalize_whitespace(str(row["preview_type"] or "")).lower() == "native": continue rel_preview_path = normalize_whitespace(str(row["rel_preview_path"] or "")) if rel_preview_path and not (paths["state_dir"] / rel_preview_path).exists(): missing_document_ids.add(int(row["document_id"])) return [ document_id for document_id in normalized_document_ids if document_id in missing_document_ids ] def load_preview_refresh_document_rows( connection: sqlite3.Connection, document_ids: list[int], ) -> list[sqlite3.Row]: normalized_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) if not normalized_document_ids: return [] return connection.execute( f""" SELECT d.id, d.rel_path, d.content_type, d.source_kind, d.conversation_id, d.production_id, d.source_text_revision_id, d.active_search_text_revision_id, active_tr.storage_rel_path AS preview_text_storage_rel_path FROM documents d LEFT JOIN text_revisions active_tr ON active_tr.id = COALESCE(d.active_search_text_revision_id, d.source_text_revision_id) WHERE d.id IN ({", ".join("?" for _ in normalized_document_ids)}) AND d.lifecycle_status NOT IN ('missing', 'deleted') ORDER BY d.id ASC """, tuple(normalized_document_ids), ).fetchall() def refresh_source_backed_document_preview( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, ) -> dict[str, object]: document_id = int(row["id"]) rel_path = str(row["rel_path"] or "") if is_internal_rel_path(rel_path): return {"status": "skipped", "reason": "internal_document"} if document_content_type_is_chat(row["content_type"]): return regenerate_chat_preview_for_document( connection, paths, document_id=document_id, ) if ( document_content_type_is_email(row["content_type"]) and row["conversation_id"] is None and row["production_id"] is None and text_revision_pointers_differ( row["source_text_revision_id"], row["active_search_text_revision_id"], ) ): return regenerate_email_preview_for_document( connection, paths, document_id=document_id, ) source_kind = normalize_whitespace(str(row["source_kind"] or FILESYSTEM_SOURCE_KIND)).lower() if source_kind != FILESYSTEM_SOURCE_KIND: return {"status": "skipped", "reason": "source_kind_requires_ingest"} source_path = document_absolute_path(paths, rel_path) if not source_path.exists(): return {"status": "skipped", "reason": "source_file_missing"} try: extracted = extract_document(source_path, include_attachments=False) except RetrieverError as exc: return {"status": "skipped", "reason": "extractor_error", "error": str(exc)} preview_artifacts = list(extracted.get("preview_artifacts", [])) if not preview_artifacts: return {"status": "skipped", "reason": "no_generated_preview_artifacts"} previous_preview_paths = [ str(preview_row["rel_preview_path"]) for preview_row in connection.execute( "SELECT rel_preview_path FROM document_previews WHERE document_id = ?", (document_id,), ).fetchall() ] preview_rows = write_preview_artifacts(paths, rel_path, preview_artifacts) replace_document_preview_rows(connection, document_id, preview_rows) cleanup_unreferenced_preview_files(paths, connection, previous_preview_paths) sync_document_attachment_preview_links(connection, paths, document_id) return {"status": "ok", "preview_rows": len(preview_rows)} def refresh_stored_state_document_preview( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, ) -> dict[str, object]: document_id = int(row["id"]) if document_content_type_is_chat(row["content_type"]): return regenerate_chat_preview_for_document( connection, paths, document_id=document_id, ) if ( document_content_type_is_email(row["content_type"]) and row["conversation_id"] is None and row["production_id"] is None ): return regenerate_email_preview_for_document( connection, paths, document_id=document_id, ) if row["production_id"] is not None: text_content = load_document_preview_text( connection, paths, document_id=document_id, storage_rel_path=row["preview_text_storage_rel_path"], ) return regenerate_production_preview_for_document( connection, paths, document_id=document_id, text_content=text_content, ) return {"status": "skipped", "reason": "requires_from_source"} def refresh_document_previews( connection: sqlite3.Connection, paths: dict[str, Path], document_ids: list[int], *, from_source: bool, ) -> dict[str, object]: refreshed = 0 skipped = 0 skip_reasons: dict[str, int] = defaultdict(int) rows = load_preview_refresh_document_rows(connection, document_ids) connection.execute("BEGIN") try: for row in rows: if row["conversation_id"] is not None and row["production_id"] is None: result = {"status": "skipped", "reason": "conversation_scope"} elif from_source: result = refresh_source_backed_document_preview(connection, paths, row) else: result = refresh_stored_state_document_preview(connection, paths, row) if result.get("status") == "ok": refreshed += 1 else: skipped += 1 reason = normalize_whitespace(str(result.get("reason") or "unknown")) skip_reasons[reason] = skip_reasons.get(reason, 0) + 1 connection.commit() except Exception: connection.rollback() raise return { "refreshed_documents": refreshed, "skipped_documents": skipped, "document_skip_reasons": dict(sorted(skip_reasons.items())), } def refresh_generated_previews( root: Path, *, scope: str = "conversations", conversation_ids: list[int] | None = None, document_ids: list[int] | None = None, dataset_id: int | None = None, dataset_name: str | None = None, missing_only: bool = False, from_source: bool = False, ) -> dict[str, object]: normalized_scope = normalize_whitespace(str(scope or "conversations")).lower() if normalized_scope not in {"conversations", "documents", "all"}: raise RetrieverError("refresh-previews --scope must be one of: conversations, documents, all.") refresh_conversation_scope = normalized_scope in {"conversations", "all"} refresh_document_scope = normalized_scope in {"documents", "all"} paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) dataset_summary: dict[str, object] | None = None target_conversation_ids: list[int] = [] candidate_conversation_ids: list[int] = [] refreshed_conversations = 0 empty_dirs_pruned = 0 if refresh_conversation_scope: conversation_dataset_summary: dict[str, object] | None target_conversation_ids, conversation_dataset_summary = resolve_conversation_preview_refresh_ids( connection, conversation_ids=conversation_ids, document_ids=document_ids, dataset_id=dataset_id, dataset_name=dataset_name, ignore_documents_without_conversation=normalized_scope == "all", ) if conversation_dataset_summary is not None: dataset_summary = conversation_dataset_summary candidate_conversation_ids = list(target_conversation_ids) if missing_only: target_conversation_ids = filter_conversation_ids_with_missing_preview_artifacts( connection, paths, target_conversation_ids, ) connection.execute("BEGIN") try: refreshed_conversations = refresh_conversation_previews(connection, paths, target_conversation_ids) connection.commit() except Exception: connection.rollback() raise empty_dirs_pruned = prune_empty_conversation_preview_dirs(paths) target_document_ids: list[int] = [] candidate_document_ids: list[int] = [] document_result: dict[str, object] = { "refreshed_documents": 0, "skipped_documents": 0, "document_skip_reasons": {}, } if refresh_document_scope: document_dataset_summary: dict[str, object] | None target_document_ids, document_dataset_summary = resolve_document_preview_refresh_ids( connection, conversation_ids=conversation_ids, document_ids=document_ids, dataset_id=dataset_id, dataset_name=dataset_name, ) if dataset_summary is None and document_dataset_summary is not None: dataset_summary = document_dataset_summary candidate_document_ids = list(target_document_ids) if missing_only: target_document_ids = filter_document_ids_with_missing_preview_artifacts( connection, paths, target_document_ids, ) document_result = refresh_document_previews( connection, paths, target_document_ids, from_source=from_source, ) result: dict[str, object] = { "status": "ok", "scope": normalized_scope, "missing_only": bool(missing_only), "from_source_requested": bool(from_source), "refreshed_conversations": int(refreshed_conversations), **document_result, "empty_conversation_preview_dirs_pruned": int(empty_dirs_pruned), } if from_source: result["from_source_mode"] = ( "enabled_for_document_previews" if refresh_document_scope else "not_applicable_for_conversation_previews" ) if refresh_conversation_scope: result["from_source_note"] = ( "Conversation previews are regenerated from stored document text and metadata; " "source reparse only applies to document preview scope." ) if missing_only: result["candidate_conversations"] = len(candidate_conversation_ids) result["candidate_documents"] = len(candidate_document_ids) if conversation_ids is None and document_ids is None and dataset_id is None and dataset_name is None: if normalized_scope == "conversations": result["target_scope"] = "all_active_conversations" elif normalized_scope == "documents": result["target_scope"] = "all_active_documents" else: result["target_scope"] = "all_active_documents_and_conversations" else: if refresh_conversation_scope: result["target_conversation_ids"] = target_conversation_ids if refresh_document_scope: result["target_document_ids"] = target_document_ids if document_ids: result["requested_document_ids"] = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) if conversation_ids: result["requested_conversation_ids"] = sorted( dict.fromkeys(int(conversation_id) for conversation_id in conversation_ids) ) if dataset_summary is not None: result["dataset"] = dataset_summary return result finally: connection.close() def rebuild_conversations( root: Path, *, conversation_ids: list[int] | None = None, document_ids: list[int] | None = None, dataset_id: int | None = None, dataset_name: str | None = None, batch_size: int = 50, ) -> dict[str, object]: normalized_batch_size = max(1, min(int(batch_size or 50), 1000)) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) raise_if_ingest_v2_active(connection, root, command_name="rebuild-conversations") connection.execute("BEGIN") try: assignment = assign_supported_conversations(connection) connection.commit() except Exception: connection.rollback() raise target_conversation_ids, dataset_summary = resolve_conversation_preview_refresh_ids( connection, conversation_ids=conversation_ids, document_ids=document_ids, dataset_id=dataset_id, dataset_name=dataset_name, ) refreshed = 0 for offset in range(0, len(target_conversation_ids), normalized_batch_size): batch_conversation_ids = target_conversation_ids[offset : offset + normalized_batch_size] connection.execute("BEGIN") try: refreshed += refresh_conversation_previews(connection, paths, batch_conversation_ids) connection.commit() except Exception: connection.rollback() raise empty_dirs_pruned = prune_empty_conversation_preview_dirs(paths) result: dict[str, object] = { "status": "ok", "assignment": {key: int(value) for key, value in dict(assignment).items()}, "target_conversations": len(target_conversation_ids), "refreshed_conversations": int(refreshed), "batch_size": normalized_batch_size, "empty_conversation_preview_dirs_pruned": int(empty_dirs_pruned), } if conversation_ids is None and document_ids is None and dataset_id is None and dataset_name is None: result["target_scope"] = "all_active_conversations" else: result["target_conversation_ids"] = target_conversation_ids if document_ids: result["requested_document_ids"] = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) if conversation_ids: result["requested_conversation_ids"] = sorted( dict.fromkeys(int(conversation_id) for conversation_id in conversation_ids) ) if dataset_summary is not None: result["dataset"] = dataset_summary return result finally: connection.close() def merge_into_conversation(root: Path, document_id: int, target_document_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) source_root_row = get_document_family_root_row_for_assignment(connection, document_id) target_root_row = get_document_family_root_row_for_assignment(connection, target_document_id) source_category = ensure_document_supports_manual_conversation_assignment(source_root_row) target_category = ensure_document_supports_manual_conversation_assignment(target_root_row) if source_category != target_category: raise RetrieverError("Source and target documents must belong to the same conversation-compatible category.") target_conversation_id = ( int(target_root_row["conversation_id"]) if target_root_row["conversation_id"] is not None else None ) if target_conversation_id is None: raise RetrieverError(f"Target document {int(target_root_row['id'])} does not belong to a conversation.") connection.execute("BEGIN") try: connection.execute( """ UPDATE documents SET conversation_id = ?, conversation_assignment_mode = ?, updated_at = ? WHERE id = ? """, ( target_conversation_id, CONVERSATION_ASSIGNMENT_MODE_MANUAL, utc_now(), int(source_root_row["id"]), ), ) assignment = reassign_conversations_and_refresh_previews(connection, paths) updated_row = get_document_row_for_conversation_assignment(connection, int(source_root_row["id"])) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "document_id": document_id, "root_document_id": int(source_root_row["id"]), "target_document_id": target_document_id, "target_root_document_id": int(target_root_row["id"]), "conversation_id": int(updated_row["conversation_id"]) if updated_row["conversation_id"] is not None else None, "conversation_assignment_mode": effective_conversation_assignment_mode( updated_row["conversation_assignment_mode"] ), "assignment_summary": assignment, } finally: connection.close() def split_from_conversation(root: Path, document_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) root_row = get_document_family_root_row_for_assignment(connection, document_id) ensure_document_supports_manual_conversation_assignment(root_row) connection.execute("BEGIN") try: singleton_conversation_id = create_manual_singleton_conversation(connection, root_row) connection.execute( """ UPDATE documents SET conversation_id = ?, conversation_assignment_mode = ?, updated_at = ? WHERE id = ? """, ( singleton_conversation_id, CONVERSATION_ASSIGNMENT_MODE_MANUAL, utc_now(), int(root_row["id"]), ), ) assignment = reassign_conversations_and_refresh_previews(connection, paths) updated_row = get_document_row_for_conversation_assignment(connection, int(root_row["id"])) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "document_id": document_id, "root_document_id": int(root_row["id"]), "conversation_id": int(updated_row["conversation_id"]) if updated_row["conversation_id"] is not None else None, "conversation_assignment_mode": effective_conversation_assignment_mode( updated_row["conversation_assignment_mode"] ), "assignment_summary": assignment, } finally: connection.close() def clear_conversation_assignment(root: Path, document_id: int) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) root_row = get_document_family_root_row_for_assignment(connection, document_id) ensure_document_supports_manual_conversation_assignment(root_row) connection.execute("BEGIN") try: connection.execute( """ UPDATE documents SET conversation_assignment_mode = ?, updated_at = ? WHERE id = ? """, ( CONVERSATION_ASSIGNMENT_MODE_AUTO, utc_now(), int(root_row["id"]), ), ) assignment = reassign_conversations_and_refresh_previews(connection, paths) updated_row = get_document_row_for_conversation_assignment(connection, int(root_row["id"])) connection.commit() except Exception: connection.rollback() raise return { "status": "ok", "document_id": document_id, "root_document_id": int(root_row["id"]), "conversation_id": int(updated_row["conversation_id"]) if updated_row["conversation_id"] is not None else None, "conversation_assignment_mode": effective_conversation_assignment_mode( updated_row["conversation_assignment_mode"] ), "assignment_summary": assignment, } finally: connection.close() def delete_docs( root: Path, *, document_ids: list[int] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, dataset_names: list[str] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, path_prefixes: list[str] | None = None, dry_run: bool = False, confirm: bool = False, ) -> dict[str, object]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) normalized_path_prefixes = normalize_delete_path_prefixes(path_prefixes) selector_inputs_present = bool( query.strip() or raw_bates or raw_filters or dataset_names or from_run_id is not None or normalized_path_prefixes ) if normalized_document_ids and (selector_inputs_present or select_from_scope): raise RetrieverError("delete-docs accepts either --doc-id selectors or query/filter/scope selectors, not both.") if dry_run and confirm: raise RetrieverError("Choose either --dry-run or --confirm, not both.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) effective_raw_filters = list(raw_filters or []) path_expression = path_prefix_scope_expression(normalized_path_prefixes) if path_expression: effective_raw_filters.append([path_expression]) occurrence_scoped = bool(normalized_path_prefixes) or raw_filters_include_occurrence_scope( connection, effective_raw_filters, ) if normalized_document_ids: target_rows = fetch_deletable_document_rows_by_ids(connection, normalized_document_ids) target_document_ids = [int(row["id"]) for row in target_rows] selector_payload: dict[str, object] = { "mode": "document_ids", "document_ids": target_document_ids, } selected_from_scope = False else: selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=effective_raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if not scope_run_selector_has_inputs(selector): raise RetrieverError( "No document selection active. Provide --doc-id, --path, or at least one of --keyword, --filter, --bates, --dataset, --from-run-id, or --select-from-scope." ) if occurrence_scoped: occurrence_document_ids = document_ids_matching_occurrence_filters(connection, effective_raw_filters) base_selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=document_only_raw_filters(connection, effective_raw_filters), dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if scope_run_selector_has_inputs(base_selector): base_document_ids, _, _ = resolve_seed_documents_for_scope_selector(connection, base_selector) allowed_document_ids = {int(document_id) for document_id in base_document_ids} scope_document_ids = [ document_id for document_id in occurrence_document_ids if document_id in allowed_document_ids ] else: scope_document_ids = occurrence_document_ids else: scope_document_ids, _, _ = resolve_seed_documents_for_scope_selector(connection, selector) target_rows = fetch_deletable_document_rows_by_ids(connection, scope_document_ids) target_document_ids = [int(row["id"]) for row in target_rows] selector_payload = { "mode": "scope_search", "scope": selector, } selected_from_scope = bool(select_from_scope) occurrence_scope_clauses: list[str] = [] occurrence_scope_params: list[object] = [] if occurrence_scoped: occurrence_scope_clauses, occurrence_scope_params = build_occurrence_scope_filters( connection, effective_raw_filters, ) root_occurrence_ids: set[int] = set() direct_occurrence_ids: set[int] = set() preview_refresh_document_ids: set[int] = set() for row in target_rows: document_id = int(row["id"]) if occurrence_scoped: occurrence_rows = scoped_active_occurrence_rows_for_document( connection, document_id, occurrence_scope_clauses, occurrence_scope_params, ) else: occurrence_rows = active_occurrence_rows_for_document(connection, document_id) if not occurrence_rows: continue if row["parent_document_id"] is None: root_occurrence_ids.update(int(occurrence_row["id"]) for occurrence_row in occurrence_rows) preview_refresh_document_ids.add(document_id) else: direct_occurrence_ids.update(int(occurrence_row["id"]) for occurrence_row in occurrence_rows) preview_refresh_document_ids.add(int(row["parent_document_id"])) targeted_occurrence_rows = collect_occurrence_rows_for_delete_plan( connection, root_occurrence_ids=root_occurrence_ids, direct_occurrence_ids=direct_occurrence_ids, ) targeted_occurrence_ids = [int(row["id"]) for row in targeted_occurrence_rows] affected_document_ids = sorted({int(row["document_id"]) for row in targeted_occurrence_rows}) preview_payload = { "selector": selector_payload, "selected_from_scope": selected_from_scope, "occurrence_scoped": occurrence_scoped, "path_prefixes": normalized_path_prefixes, "matched_document_count": len(target_document_ids), "document_ids": target_document_ids, "affected_document_ids": affected_document_ids, "would_delete_occurrences": len(targeted_occurrence_ids), "would_touch_documents": len(affected_document_ids), "sample": sample_documents_for_delete(connection, target_document_ids), } if not targeted_occurrence_ids: return { "status": "ok", **preview_payload, "deleted_occurrences": 0, "deleted_document_ids": [], "deleted_documents": 0, "retained_document_ids": [], "retained_documents": 0, "dataset_memberships_removed": 0, "attachment_preview_updates": 0, "assignment_summary": {"documents_assigned": 0, "conversations_created": 0}, } if dry_run: return {"status": "ok", "dry_run": True, **preview_payload} if not confirm: return {"status": "confirm_required", **preview_payload} connection.execute("BEGIN") try: deleted_at = utc_now() placeholders = ", ".join("?" for _ in targeted_occurrence_ids) deleted_occurrence_cursor = connection.execute( f""" UPDATE document_occurrences SET lifecycle_status = 'deleted', updated_at = ? WHERE lifecycle_status != 'deleted' AND id IN ({placeholders}) """, [deleted_at, *targeted_occurrence_ids], ) for document_id in affected_document_ids: refresh_source_backed_dataset_memberships_for_document(connection, document_id) refresh_document_from_occurrences(connection, document_id) deleted_document_ids = delete_documents_with_no_active_occurrences( connection, paths, affected_document_ids, deleted_at=deleted_at, ) dataset_memberships_removed = delete_dataset_memberships_for_documents(connection, deleted_document_ids) assignment_summary = reassign_conversations_and_refresh_previews(connection, paths) attachment_preview_updates = 0 for preview_document_id in sorted(preview_refresh_document_ids): preview_document_row = connection.execute( """ SELECT lifecycle_status FROM documents WHERE id = ? """, (preview_document_id,), ).fetchone() if preview_document_row is None or preview_document_row["lifecycle_status"] in {"missing", "deleted"}: continue attachment_preview_updates += sync_document_attachment_preview_links( connection, paths, preview_document_id, ) connection.commit() except Exception: connection.rollback() raise deleted_document_id_set = {int(document_id) for document_id in deleted_document_ids} retained_document_ids = [ document_id for document_id in affected_document_ids if document_id not in deleted_document_id_set ] return { "status": "ok", **preview_payload, "deleted_occurrences": int(deleted_occurrence_cursor.rowcount or 0), "deleted_document_ids": deleted_document_ids, "deleted_documents": len(deleted_document_ids), "retained_document_ids": retained_document_ids, "retained_documents": len(retained_document_ids), "dataset_memberships_removed": dataset_memberships_removed, "attachment_preview_updates": attachment_preview_updates, "assignment_summary": assignment_summary, } finally: connection.close() def fill_field( root: Path, *, field_name: str, value: str | None = None, clear: bool = False, document_ids: list[int] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, dataset_names: list[str] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, dry_run: bool = False, confirm: bool = False, ) -> dict[str, object]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) selector_inputs_present = bool(query.strip() or raw_bates or raw_filters or dataset_names or from_run_id is not None) if normalized_document_ids and (selector_inputs_present or select_from_scope): raise RetrieverError("fill-field accepts either --doc-id selectors or query/filter/scope selectors, not both.") if clear and value is not None: raise RetrieverError("Choose either --value or --clear, not both.") if not clear and value is None: raise RetrieverError("Provide --value or --clear.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) field_def = ensure_fill_target_field_definition(connection, field_name) column_name = str(field_def["field_name"]) typed_value = None if clear else value_from_type(str(field_def["field_type"]), value) if normalized_document_ids: target_rows = fetch_visible_document_rows_by_ids(connection, normalized_document_ids) target_document_ids = [int(row["id"]) for row in target_rows] selector_payload: dict[str, object] = { "mode": "document_ids", "document_ids": target_document_ids, } selected_from_scope = False else: selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if not scope_run_selector_has_inputs(selector): raise RetrieverError(build_no_document_selection_error()) target_document_ids, _, _ = resolve_seed_documents_for_scope_selector(connection, selector) selector_payload = { "mode": "scope_search", "scope": selector, } selected_from_scope = bool(select_from_scope) single_explicit_document = bool(normalized_document_ids) and len(target_document_ids) == 1 preview_payload = { "field_name": column_name, "field_type": str(field_def["field_type"]), "value": typed_value, "clear": clear, "selector": selector_payload, "selected_from_scope": selected_from_scope, "would_write": len(target_document_ids), "document_ids": target_document_ids, "sample": sample_documents_for_fill(connection, target_document_ids, column_name), } if dry_run: return {"status": "ok", "dry_run": True, **preview_payload} if target_document_ids and not single_explicit_document and not confirm: return {"status": "confirm_required", **preview_payload} if not target_document_ids: return { "status": "ok", "field_name": column_name, "field_type": str(field_def["field_type"]), "value": typed_value, "selector": selector_payload, "selected_from_scope": selected_from_scope, "written": 0, "skipped": 0, "failed": 0, "document_ids": [], "sample": [], } placeholders = ", ".join("?" for _ in target_document_ids) rows = connection.execute( f""" SELECT id, {quote_identifier(column_name)} AS field_value, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} AS locks_json FROM documents WHERE id IN ({placeholders}) """, tuple(target_document_ids), ).fetchall() rows_by_id = {int(row["id"]): row for row in rows} connection.execute("BEGIN") try: written = 0 skipped = 0 updated_document_ids: list[int] = [] for document_id in target_document_ids: row = rows_by_id.get(document_id) if row is None: raise RetrieverError(f"Unknown document id: {document_id}") locks = normalize_string_list(row["locks_json"]) next_locks = list(locks) if column_name not in next_locks: next_locks.append(column_name) if row["field_value"] == typed_value and next_locks == locks: skipped += 1 continue connection.execute( f""" UPDATE documents SET {quote_identifier(column_name)} = ?, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} = ?, updated_at = ? WHERE id = ? """, (typed_value, json.dumps(next_locks), utc_now(), document_id), ) updated_document_ids.append(document_id) written += 1 if column_name in {"author", "participants", "recipients", "subject", "title"}: for document_id in updated_document_ids: refresh_documents_fts_row(connection, document_id) connection.commit() return { "status": "ok", "field_name": column_name, "field_type": str(field_def["field_type"]), "value": typed_value, "selector": selector_payload, "selected_from_scope": selected_from_scope, "written": written, "skipped": skipped, "failed": 0, "document_ids": target_document_ids, "sample": sample_documents_for_fill(connection, target_document_ids, column_name), } except Exception: connection.rollback() raise finally: connection.close() def set_field(root: Path, document_id_or_ids: int | list[int], field_name: str, value: str | None) -> dict[str, object]: if isinstance(document_id_or_ids, int): target_document_ids = [document_id_or_ids] else: target_document_ids = list(dict.fromkeys(int(document_id) for document_id in document_id_or_ids)) if not target_document_ids: raise RetrieverError("set-field requires at least one document id.") payload = fill_field( root, field_name=field_name, value=value, clear=value is None, document_ids=target_document_ids, confirm=True, ) if len(target_document_ids) != 1: return { "status": str(payload.get("status") or "ok"), "field_name": str(payload.get("field_name") or field_name), "field_type": str(payload.get("field_type") or ""), "value": payload.get("value"), "written": int(payload.get("written") or 0), "skipped": int(payload.get("skipped") or 0), "failed": int(payload.get("failed") or 0), "document_ids": [int(document_id) for document_id in payload.get("document_ids") or target_document_ids], "sample": payload.get("sample") or [], } document_id = target_document_ids[0] paths = workspace_paths(root) connection = connect_db(paths["db_path"]) try: field_def = ensure_fill_target_field_definition(connection, field_name) row = connection.execute( f""" SELECT {quote_identifier(field_def['field_name'])} AS field_value, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} AS locks_json FROM documents WHERE id = ? """, (document_id,), ).fetchone() locks = normalize_string_list(row["locks_json"] if row is not None else None) return { "status": str(payload.get("status") or "ok"), "document_id": document_id, "field_name": str(field_def["field_name"]), "field_type": str(field_def["field_type"]), "value": row["field_value"] if row is not None else None, "manual_field_locks": locks, } finally: connection.close() TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN = { "author": "extracted_author", "participants": "extracted_participants", "recipients": "extracted_recipients", "date_created": "extracted_doc_authored_at", "subject": "extracted_subject", "title": "extracted_title", } DEFAULT_TEXT_DERIVED_METADATA_FIELDS = ("author", "participants", "date_created") ENTITY_REBUILD_TEXT_METADATA_FIELDS = frozenset({"author", "participants", "recipients"}) def normalize_text_derived_metadata_fields(raw_fields: list[str] | None = None) -> list[str]: normalized_fields: list[str] = [] for raw_field in list(raw_fields or DEFAULT_TEXT_DERIVED_METADATA_FIELDS): normalized_field = normalize_whitespace(str(raw_field or "")).lower() if not normalized_field: continue if normalized_field not in TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN: supported = ", ".join(sorted(TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN)) raise RetrieverError( f"Unsupported text-derived metadata field {raw_field!r}. Supported values: {supported}." ) if normalized_field not in normalized_fields: normalized_fields.append(normalized_field) if not normalized_fields: raise RetrieverError("At least one metadata field must be selected.") return normalized_fields def load_active_text_for_document( connection: sqlite3.Connection, paths: dict[str, Path], document_row: sqlite3.Row, ) -> tuple[str | None, int | None]: revision_id = ( int(document_row["active_search_text_revision_id"]) if document_row["active_search_text_revision_id"] is not None else ( int(document_row["source_text_revision_id"]) if document_row["source_text_revision_id"] is not None else None ) ) if revision_id is not None: revision_row = require_text_revision_row_by_id(connection, revision_id) text_content = read_text_revision_body(paths, revision_row["storage_rel_path"]) if text_content is not None: return text_content, revision_id chunk_rows = document_chunk_rows(connection, int(document_row["id"])) if not chunk_rows: return None, revision_id return "\n\n".join(str(chunk_row["text_content"] or "") for chunk_row in chunk_rows), revision_id def derive_text_metadata_from_active_text(text_content: str) -> dict[str, str | None]: metadata = extract_email_like_headers(text_content) if not metadata: return {} chain_participants = extract_email_chain_participants( text_content, [ metadata.get("author"), metadata.get("participants"), metadata.get("recipients"), ], ) if chain_participants: metadata["participants"] = chain_participants return metadata def build_rebuild_entities_followup_commands(root: Path, document_ids: list[int]) -> list[str]: normalized_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) if not normalized_document_ids: return [] root_arg = shlex.quote(str(root)) doc_args = " ".join(f"--doc-id {document_id}" for document_id in normalized_document_ids) return [f"rebuild-entities {root_arg} {doc_args}"] def build_refresh_text_derived_metadata_followup_commands( root: Path, document_ids: list[int], *, field_names: list[str] | None = None, ) -> list[str]: normalized_document_ids = sorted(dict.fromkeys(int(document_id) for document_id in document_ids)) if not normalized_document_ids: return [] normalized_fields = normalize_text_derived_metadata_fields(field_names) root_arg = shlex.quote(str(root)) doc_args = " ".join(f"--doc-id {document_id}" for document_id in normalized_document_ids) field_args = ( " ".join(f"--field {field_name}" for field_name in normalized_fields) if tuple(normalized_fields) != DEFAULT_TEXT_DERIVED_METADATA_FIELDS else "" ) args = " ".join(part for part in [field_args, doc_args] if part) return [f"refresh-text-derived-metadata {root_arg} {args}".strip()] def refresh_text_derived_metadata_for_document_rows( connection: sqlite3.Connection, paths: dict[str, Path], *, document_rows: list[sqlite3.Row], field_names: list[str], ) -> dict[str, object]: normalized_field_names = normalize_text_derived_metadata_fields(field_names) updated_documents = 0 updated_fields = 0 skipped_reasons: dict[str, int] = defaultdict(int) results: list[dict[str, object]] = [] entity_rebuild_candidate_document_ids: list[int] = [] now = utc_now() for document_row in document_rows: document_id = int(document_row["id"]) before = { field_name: document_row[field_name] for field_name in TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN } manual_locks = set(normalize_string_list(document_row[MANUAL_FIELD_LOCKS_COLUMN])) text_content, revision_id = load_active_text_for_document(connection, paths, document_row) if not normalize_whitespace(text_content or ""): skipped_reasons["no_active_text"] = skipped_reasons.get("no_active_text", 0) + 1 results.append( { "document_id": document_id, "status": "skipped", "reason": "no_active_text", "active_text_revision_id": revision_id, "field_names": normalized_field_names, "before": before, "after": dict(before), "updated_fields": [], "locked_fields": [], "derived": {}, } ) continue derived = derive_text_metadata_from_active_text(str(text_content or "")) if not derived: skipped_reasons["no_email_headers"] = skipped_reasons.get("no_email_headers", 0) + 1 results.append( { "document_id": document_id, "status": "skipped", "reason": "no_email_headers", "active_text_revision_id": revision_id, "field_names": normalized_field_names, "before": before, "after": dict(before), "updated_fields": [], "locked_fields": [], "derived": {}, } ) continue occurrence_rows = active_occurrence_rows_for_document(connection, document_id) if not occurrence_rows: skipped_reasons["no_active_occurrences"] = skipped_reasons.get("no_active_occurrences", 0) + 1 results.append( { "document_id": document_id, "status": "skipped", "reason": "no_active_occurrences", "active_text_revision_id": revision_id, "field_names": normalized_field_names, "before": before, "after": dict(before), "updated_fields": [], "locked_fields": [], "derived": { field_name: derived.get(field_name) for field_name in normalized_field_names if derived.get(field_name) is not None }, } ) continue selected_derived_values = { field_name: derived.get(field_name) for field_name in normalized_field_names if derived.get(field_name) is not None } if not selected_derived_values: skipped_reasons["no_selected_fields"] = skipped_reasons.get("no_selected_fields", 0) + 1 results.append( { "document_id": document_id, "status": "skipped", "reason": "no_selected_fields", "active_text_revision_id": revision_id, "field_names": normalized_field_names, "before": before, "after": dict(before), "updated_fields": [], "locked_fields": [], "derived": {}, } ) continue for occurrence_row in occurrence_rows: assignments: list[str] = [] params: list[object] = [] for field_name, derived_value in selected_derived_values.items(): column_name = TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN[field_name] if occurrence_row[column_name] == derived_value: continue assignments.append(f"{quote_identifier(column_name)} = ?") params.append(derived_value) if not assignments: continue assignments.append("updated_at = ?") params.append(now) params.append(int(occurrence_row["id"])) connection.execute( f""" UPDATE document_occurrences SET {', '.join(assignments)} WHERE id = ? """, params, ) refresh_document_from_occurrences(connection, document_id) updated_row = connection.execute( f""" SELECT id, {quote_identifier(MANUAL_FIELD_LOCKS_COLUMN)} AS locks_json, {", ".join(quote_identifier(field_name) for field_name in TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN)} FROM documents WHERE id = ? """, (document_id,), ).fetchone() assert updated_row is not None after = { field_name: updated_row[field_name] for field_name in TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN } changed_field_names = [ field_name for field_name in normalized_field_names if before.get(field_name) != after.get(field_name) ] locked_fields = [ field_name for field_name in normalized_field_names if field_name in manual_locks and selected_derived_values.get(field_name) is not None and before.get(field_name) != selected_derived_values.get(field_name) ] if changed_field_names: updated_documents += 1 updated_fields += len(changed_field_names) if set(changed_field_names) & ENTITY_REBUILD_TEXT_METADATA_FIELDS: entity_rebuild_candidate_document_ids.append(document_id) results.append( { "document_id": document_id, "status": "updated" if changed_field_names else "unchanged", "reason": None, "active_text_revision_id": revision_id, "field_names": normalized_field_names, "before": before, "after": after, "updated_fields": changed_field_names, "locked_fields": locked_fields, "derived": selected_derived_values, } ) return { "status": "ok", "field_names": normalized_field_names, "processed_documents": len(document_rows), "updated_documents": updated_documents, "updated_fields": updated_fields, "skipped_documents": sum(1 for item in results if item["status"] == "skipped"), "skipped_reasons": dict(sorted(skipped_reasons.items())), "results": results, "entity_rebuild_candidate_document_ids": sorted(dict.fromkeys(entity_rebuild_candidate_document_ids)), } def refresh_text_derived_metadata( root: Path, *, document_ids: list[int] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, dataset_names: list[str] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, field_names: list[str] | None = None, ) -> dict[str, object]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) normalized_field_names = normalize_text_derived_metadata_fields(field_names) selector_inputs_present = bool(query.strip() or raw_bates or raw_filters or dataset_names or from_run_id is not None) if normalized_document_ids and (selector_inputs_present or select_from_scope): raise RetrieverError( "refresh-text-derived-metadata accepts either --doc-id selectors or query/filter/scope selectors, not both." ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) if normalized_document_ids: target_rows = fetch_visible_document_rows_by_ids(connection, normalized_document_ids) selector_payload: dict[str, object] = { "mode": "document_ids", "document_ids": [int(row["id"]) for row in target_rows], } selected_from_scope = False else: selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if not scope_run_selector_has_inputs(selector): raise RetrieverError(build_no_document_selection_error()) target_document_ids, _, _ = resolve_seed_documents_for_scope_selector(connection, selector) target_rows = fetch_visible_document_rows_by_ids(connection, target_document_ids) selector_payload = { "mode": "scope_search", "scope": selector, } selected_from_scope = bool(select_from_scope) connection.execute("BEGIN") try: payload = refresh_text_derived_metadata_for_document_rows( connection, paths, document_rows=target_rows, field_names=normalized_field_names, ) connection.commit() except Exception: connection.rollback() raise payload["selector"] = selector_payload payload["selected_from_scope"] = selected_from_scope payload["document_ids"] = [int(row["id"]) for row in target_rows] payload["next_recommended_commands"] = build_rebuild_entities_followup_commands( root, list(payload.get("entity_rebuild_candidate_document_ids") or []), ) return payload finally: connection.close() OCCURRENCE_FILTER_FIELDS = { "begin_attachment", "begin_bates", "custodian", "end_attachment", "end_bates", "file_hash", "file_name", "file_size", "file_type", "ingested_at", "last_seen_at", "lifecycle_status", "production_id", "rel_path", "source_folder_path", "source_item_id", "source_kind", "source_rel_path", "text_status", "updated_at", } ENTITY_ROLE_FILTER_FIELDS = { "author": "author", "participant": "participant", "participants": "participant", "recipient": "recipient", "recipients": "recipient", "custodian": "custodian", } ENTITY_ROLE_ID_FILTER_FIELDS = { "author_entity_id": "author", "participant_entity_id": "participant", "participants_entity_id": "participant", "recipient_entity_id": "recipient", "recipients_entity_id": "recipient", "custodian_entity_id": "custodian", } RAW_ENTITY_METADATA_FILTER_FIELDS = { "raw_author": "extracted_author", "raw_participant": "extracted_participants", "raw_participants": "extracted_participants", "raw_recipient": "extracted_recipients", "raw_recipients": "extracted_recipients", "raw_custodian": "custodian", } ROLE_OCCURRENCE_METADATA_COLUMNS = { "author": "extracted_author", "participant": "extracted_participants", "recipient": "extracted_recipients", "custodian": "custodian", } def build_scalar_filter_clause( column_expr: str, field_type: str, operator: str, value: str | None, ) -> tuple[str, list[object]]: if operator == "is-null": return f"{column_expr} IS NULL", [] if operator == "not-null": return f"{column_expr} IS NOT NULL", [] if operator == "contains": if field_type not in {"text", "date"}: raise RetrieverError(f"Operator 'contains' is not valid for field type '{field_type}'.") return f"LOWER(COALESCE({column_expr}, '')) LIKE LOWER(?)", [f"%{value}%"] typed_value = value_from_type(field_type if field_type in {"integer", "real", "boolean"} else "text", value) if field_type == "date": typed_value = normalize_date_field_value(str(value or "")) if typed_value is None: raise RetrieverError(f"Expected ISO date value, got {value!r}") comparators = { "eq": "=", "neq": "!=", "gt": ">", "gte": ">=", "lt": "<", "lte": "<=", } comparator = comparators[operator] if operator in {"gt", "gte", "lt", "lte"} and field_type not in {"integer", "real", "date", "text"}: raise RetrieverError(f"Operator '{operator}' is not valid for field type '{field_type}'.") return f"{column_expr} {comparator} ?", [typed_value] def build_filter_clause(alias: str, field_def: dict[str, str], operator: str, value: str | None) -> tuple[str, list[object]]: field_name = field_def["field_name"] field_type = field_def["field_type"] field_source = field_def.get("source") if field_source == "virtual": return build_virtual_filter_clause(alias, field_name, field_type, operator, value) if alias == "d" and field_name in OCCURRENCE_FILTER_FIELDS: occurrence_clause, occurrence_params = build_scalar_filter_clause( f"o.{quote_identifier(field_name)}", field_type, operator, value, ) return ( "EXISTS (" "SELECT 1 FROM document_occurrences o " f"WHERE o.document_id = {alias}.id " "AND o.lifecycle_status = 'active' " f"AND {occurrence_clause}" ")", occurrence_params, ) return build_scalar_filter_clause(f"{alias}.{quote_identifier(field_name)}", field_type, operator, value) def build_virtual_filter_clause( alias: str, field_name: str, field_type: str, operator: str, value: str | None, ) -> tuple[str, list[object]]: if field_name == "custodian": exists_expr = ( "EXISTS (" "SELECT 1 " "FROM document_occurrences o " f"WHERE o.document_id = {alias}.id " "AND o.lifecycle_status = 'active'" ) filtered_exists_expr = ( "EXISTS (" "SELECT 1 " "FROM document_occurrences o " f"WHERE o.document_id = {alias}.id " "AND o.lifecycle_status = 'active' " ) if operator == "is-null": return f"NOT {filtered_exists_expr} AND COALESCE(o.custodian, '') != '')", [] if operator == "not-null": return f"{filtered_exists_expr} AND COALESCE(o.custodian, '') != '')", [] if operator == "contains": return ( f"{filtered_exists_expr} AND LOWER(COALESCE(o.custodian, '')) LIKE LOWER(?))", [f"%{value}%"], ) if operator == "neq": return f"NOT ({filtered_exists_expr} AND COALESCE(o.custodian, '') = ?)", [value or ""] occurrence_clause, occurrence_params = build_scalar_filter_clause( "o.custodian", field_type, operator, value, ) return f"{exists_expr} AND {occurrence_clause})", occurrence_params if field_name in {"is_attachment", "has_attachments"}: if operator not in {"eq", "neq"}: raise RetrieverError(f"Virtual filter '{field_name}' only supports eq and neq.") typed_value = value_from_type("boolean", value) if field_name == "is_attachment": positive_clause = attachment_child_filter_sql(alias) else: positive_clause = ( "EXISTS (" "SELECT 1 FROM documents child " f"WHERE child.parent_document_id = {alias}.id " f"AND COALESCE(child.child_document_kind, '{CHILD_DOCUMENT_KIND_ATTACHMENT}') = '{CHILD_DOCUMENT_KIND_ATTACHMENT}' " "AND child.lifecycle_status NOT IN ('missing', 'deleted')" ")" ) positive = bool(typed_value) if operator == "neq": positive = not positive return (positive_clause if positive else f"NOT ({positive_clause})"), [] if field_name == "production_name": column_expr = ( "(SELECT p.production_name FROM productions p " f"WHERE p.id = {alias}.production_id)" ) if operator == "is-null": return f"{column_expr} IS NULL", [] if operator == "not-null": return f"{column_expr} IS NOT NULL", [] if operator == "contains": return f"LOWER(COALESCE({column_expr}, '')) LIKE LOWER(?)", [f"%{value}%"] if operator in {"eq", "neq"}: comparator = "=" if operator == "eq" else "!=" return f"COALESCE({column_expr}, '') {comparator} ?", [value or ""] raise RetrieverError(f"Virtual filter '{field_name}' does not support operator '{operator}'.") if field_name == "dataset_name": exists_expr = ( "EXISTS (" "SELECT 1 " "FROM dataset_documents dd " "JOIN datasets ds ON ds.id = dd.dataset_id " f"WHERE dd.document_id = {alias}.id" ) filtered_exists_expr = ( "EXISTS (" "SELECT 1 " "FROM dataset_documents dd " "JOIN datasets ds ON ds.id = dd.dataset_id " f"WHERE dd.document_id = {alias}.id " ) if operator == "is-null": return f"NOT {exists_expr})", [] if operator == "not-null": return f"{exists_expr})", [] if operator == "contains": return f"{filtered_exists_expr} AND LOWER(COALESCE(ds.dataset_name, '')) LIKE LOWER(?))", [f"%{value}%"] if operator in {"eq", "neq"}: positive_clause = f"{filtered_exists_expr} AND COALESCE(ds.dataset_name, '') = ?)" if operator == "eq": return positive_clause, [value or ""] return f"NOT ({positive_clause})", [value or ""] raise RetrieverError(f"Virtual filter '{field_name}' does not support operator '{operator}'.") raise RetrieverError(f"Unknown virtual filter: {field_name}") def build_search_filters( connection: sqlite3.Connection, raw_filters: object | None ) -> tuple[list[str], list[str], list[object]]: return build_sql_like_search_filters(connection, raw_filters) def base_document_search_clauses() -> list[str]: return [ "d.lifecycle_status NOT IN ('missing', 'deleted')", "EXISTS (SELECT 1 FROM dataset_documents dd WHERE dd.document_id = d.id)", ] def known_logical_field_names(connection: sqlite3.Connection) -> list[str]: names = set(BUILTIN_FIELD_TYPES) | set(VIRTUAL_FILTER_FIELD_TYPES) if table_exists(connection, "custom_fields_registry"): registry_rows = connection.execute( """ SELECT field_name FROM custom_fields_registry ORDER BY field_name ASC """ ).fetchall() document_columns = table_columns(connection, "documents") for row in registry_rows: field_name = str(row["field_name"]) if field_name in document_columns: names.add(field_name) return sorted(names) def field_name_suggestions(connection: sqlite3.Connection, field_name: str) -> list[str]: candidates = known_logical_field_names(connection) return difflib.get_close_matches(field_name, candidates, n=3, cutoff=0.45) def sql_filter_operator_names_for_field_type(field_type: str) -> list[str]: if field_type == "boolean": return ["=", "!=", "IS NULL", "IS NOT NULL"] operators = ["=", "!=", "<", "<=", ">", ">=", "IS NULL", "IS NOT NULL", "IN", "BETWEEN"] if field_type in {"text", "date"}: operators.insert(6, "LIKE") return operators def filter_error_excerpt(expression: str, position: int) -> tuple[str, str]: start = max(0, position - 60) end = min(len(expression), position + 60) excerpt = expression[start:end] caret = " " * max(0, position - start) + "^" return excerpt, caret def raise_filter_syntax_error(expression: str, position: int, message: str) -> None: excerpt, caret = filter_error_excerpt(expression, position) raise RetrieverError(f"{message} at position {position + 1}.\n{excerpt}\n{caret}") def tokenize_sql_filter_expression(expression: str) -> list[dict[str, object]]: if len(expression.encode("utf-8")) > MAX_FILTER_EXPRESSION_LENGTH: raise RetrieverError( f"Filter expressions are capped at {MAX_FILTER_EXPRESSION_LENGTH} bytes." ) tokens: list[dict[str, object]] = [] index = 0 length = len(expression) keywords = {"AND", "BETWEEN", "FALSE", "IN", "IS", "LIKE", "NOT", "NULL", "OR", "TRUE"} while index < length: char = expression[index] if char.isspace(): index += 1 continue if char in "(),": token_kind = {"(": "lparen", ")": "rparen", ",": "comma"}[char] tokens.append({"kind": token_kind, "value": char, "start": index, "end": index + 1}) index += 1 continue if expression.startswith("<=", index) or expression.startswith(">=", index) or expression.startswith("!=", index) or expression.startswith("<>", index): tokens.append({"kind": "operator", "value": expression[index:index + 2], "start": index, "end": index + 2}) index += 2 continue if char in "=<>": tokens.append({"kind": "operator", "value": char, "start": index, "end": index + 1}) index += 1 continue if char in {"'", '"'}: quote = char start = index index += 1 value_chars: list[str] = [] while index < length: current = expression[index] if current == "\\" and index + 1 < length: value_chars.append(expression[index + 1]) index += 2 continue if current == quote: if index + 1 < length and expression[index + 1] == quote: value_chars.append(quote) index += 2 continue index += 1 break value_chars.append(current) index += 1 else: raise_filter_syntax_error(expression, start, "Unterminated string literal") tokens.append( { "kind": "literal", "literal_kind": "string", "value": "".join(value_chars), "start": start, "end": index, } ) continue if char in "+-" and index + 1 < length and expression[index + 1].isdigit(): start = index index += 1 while index < length and expression[index].isdigit(): index += 1 if index < length and expression[index] == ".": index += 1 while index < length and expression[index].isdigit(): index += 1 tokens.append( { "kind": "literal", "literal_kind": "number", "value": expression[start:index], "start": start, "end": index, } ) continue if char.isdigit(): start = index while index < length and expression[index].isdigit(): index += 1 if index < length and expression[index] == ".": index += 1 while index < length and expression[index].isdigit(): index += 1 tokens.append( { "kind": "literal", "literal_kind": "number", "value": expression[start:index], "start": start, "end": index, } ) continue if char.isalpha() or char == "_": start = index index += 1 while index < length and (expression[index].isalnum() or expression[index] == "_"): index += 1 raw_value = expression[start:index] upper_value = raw_value.upper() if upper_value in keywords: if upper_value in {"TRUE", "FALSE"}: tokens.append( { "kind": "literal", "literal_kind": "boolean", "value": upper_value == "TRUE", "start": start, "end": index, } ) elif upper_value == "NULL": tokens.append( { "kind": "literal", "literal_kind": "null", "value": None, "start": start, "end": index, } ) else: tokens.append({"kind": "keyword", "value": upper_value, "start": start, "end": index}) else: tokens.append({"kind": "identifier", "value": raw_value, "start": start, "end": index}) continue raise_filter_syntax_error(expression, index, f"Unexpected character {char!r}") tokens.append({"kind": "eof", "value": "", "start": length, "end": length}) return tokens def peek_filter_token(state: dict[str, object]) -> dict[str, object]: tokens = state["tokens"] index = int(state["index"]) return tokens[index] # type: ignore[index] def consume_filter_token(state: dict[str, object]) -> dict[str, object]: token = peek_filter_token(state) state["index"] = int(state["index"]) + 1 return token def match_filter_keyword(state: dict[str, object], keyword: str) -> bool: token = peek_filter_token(state) if token["kind"] == "keyword" and token["value"] == keyword: consume_filter_token(state) return True return False def match_filter_token_kind(state: dict[str, object], kind: str, value: str | None = None) -> dict[str, object] | None: token = peek_filter_token(state) if token["kind"] != kind: return None if value is not None and token["value"] != value: return None consume_filter_token(state) return token def expect_filter_token(state: dict[str, object], kind: str, message: str, value: str | None = None) -> dict[str, object]: token = match_filter_token_kind(state, kind, value=value) if token is not None: return token next_token = peek_filter_token(state) raise_filter_syntax_error(str(state["expression"]), int(next_token["start"]), message) def resolve_sql_filter_field(connection: sqlite3.Connection, raw_field_name: str) -> dict[str, object]: try: return resolve_field_definition(connection, raw_field_name) except RetrieverError as exc: suggestions = field_name_suggestions(connection, raw_field_name) suggestion_text = f" Did you mean: {', '.join(suggestions)}?" if suggestions else "" raise RetrieverError(f"Unknown field '{raw_field_name}'.{suggestion_text}") from exc def literal_text_value(literal: dict[str, object]) -> str | None: literal_kind = literal["literal_kind"] if literal_kind == "null": return None if literal_kind == "boolean": return "true" if bool(literal["value"]) else "false" return str(literal["value"]) def coerce_sql_literal(field_type: str, literal: dict[str, object]) -> object: raw_value = literal_text_value(literal) if raw_value is None: raise RetrieverError("NULL is only valid with IS NULL / IS NOT NULL.") if literal["literal_kind"] == "boolean" and field_type not in {"boolean", "text"}: raise RetrieverError(f"Expected {field_type} value, got boolean literal.") if literal["literal_kind"] == "number" and field_type == "boolean": return value_from_type("boolean", raw_value) if field_type == "date": normalized = normalize_date_field_value(raw_value) if normalized is None: raise RetrieverError(f"Expected ISO date value, got {raw_value!r}") return normalized if field_type in {"integer", "real", "boolean"}: return value_from_type(field_type, raw_value) return raw_value def ensure_sql_filter_operator_supported(field_def: dict[str, object], operator: str) -> None: supported = sql_filter_operator_names_for_field_type(str(field_def["field_type"])) if operator not in supported: field_name = str(field_def["field_name"]) field_type = str(field_def["field_type"]) raise RetrieverError( f"Field '{field_name}' is {field_type}; supported operators: {', '.join(supported)}." ) def build_scalar_sql_filter_clause( sql_expression: str, field_def: dict[str, object], operator: str, operand: object | None, ) -> tuple[str, list[object]]: ensure_sql_filter_operator_supported(field_def, operator) field_type = str(field_def["field_type"]) if operator in {"IS NULL", "IS NOT NULL"}: return f"{sql_expression} {operator}", [] if operator == "LIKE": if field_type not in {"text", "date"}: ensure_sql_filter_operator_supported(field_def, operator) assert isinstance(operand, dict) return f"COALESCE({sql_expression}, '') LIKE ?", [str(coerce_sql_literal("text", operand))] if operator == "IN": if not isinstance(operand, list) or not operand: raise RetrieverError("IN requires at least one value.") if len(operand) > MAX_FILTER_IN_LIST_ITEMS: raise RetrieverError(f"IN (...) is capped at {MAX_FILTER_IN_LIST_ITEMS} values.") typed_values = [coerce_sql_literal(field_type, literal) for literal in operand] placeholders = ", ".join("?" for _ in typed_values) return f"{sql_expression} IN ({placeholders})", typed_values if operator == "BETWEEN": if not isinstance(operand, tuple) or len(operand) != 2: raise RetrieverError("BETWEEN requires two values.") left_value = coerce_sql_literal(field_type, operand[0]) right_value = coerce_sql_literal(field_type, operand[1]) return f"{sql_expression} BETWEEN ? AND ?", [left_value, right_value] assert isinstance(operand, dict) comparator = "!=" if operator == "<>" else operator return f"{sql_expression} {comparator} ?", [coerce_sql_literal(field_type, operand)] def build_dataset_name_sql_filter_clause( alias: str, field_def: dict[str, object], operator: str, operand: object | None, ) -> tuple[str, list[object]]: ensure_sql_filter_operator_supported(field_def, operator) exists_sql = ( "SELECT 1 " "FROM dataset_documents dd " "JOIN datasets ds ON ds.id = dd.dataset_id " f"WHERE dd.document_id = {alias}.id" ) if operator == "IS NULL": return f"NOT EXISTS ({exists_sql})", [] if operator == "IS NOT NULL": return f"EXISTS ({exists_sql})", [] if operator == "IN": assert isinstance(operand, list) if len(operand) > MAX_FILTER_IN_LIST_ITEMS: raise RetrieverError(f"IN (...) is capped at {MAX_FILTER_IN_LIST_ITEMS} values.") values = [str(coerce_sql_literal("text", literal)) for literal in operand] placeholders = ", ".join("?" for _ in values) return f"EXISTS ({exists_sql} AND ds.dataset_name IN ({placeholders}))", values if operator == "BETWEEN": assert isinstance(operand, tuple) values = [str(coerce_sql_literal("text", operand[0])), str(coerce_sql_literal("text", operand[1]))] return f"EXISTS ({exists_sql} AND ds.dataset_name BETWEEN ? AND ?)", values if operator == "LIKE": assert isinstance(operand, dict) value = str(coerce_sql_literal("text", operand)) return f"EXISTS ({exists_sql} AND COALESCE(ds.dataset_name, '') LIKE ?)", [value] assert isinstance(operand, dict) comparator = "!=" if operator == "<>" else operator value = str(coerce_sql_literal("text", operand)) if comparator == "!=": return f"NOT EXISTS ({exists_sql} AND COALESCE(ds.dataset_name, '') = ?)", [value] return f"EXISTS ({exists_sql} AND COALESCE(ds.dataset_name, '') {comparator} ?)", [value] def build_custodian_sql_filter_clause( alias: str, field_def: dict[str, object], operator: str, operand: object | None, *, occurrence_alias: str | None = None, ) -> tuple[str, list[object]]: ensure_sql_filter_operator_supported(field_def, operator) if occurrence_alias is not None: sql_expression = f"{occurrence_alias}.custodian" if operator == "IS NULL": return f"COALESCE({sql_expression}, '') = ''", [] if operator == "IS NOT NULL": return f"COALESCE({sql_expression}, '') != ''", [] if operator in {"!=", "<>"}: assert isinstance(operand, dict) value = str(coerce_sql_literal("text", operand)) return f"COALESCE({sql_expression}, '') != ?", [value] return build_scalar_sql_filter_clause(sql_expression, field_def, operator, operand) exists_sql = ( "SELECT 1 " "FROM document_occurrences o " f"WHERE o.document_id = {alias}.id " "AND o.lifecycle_status = 'active'" ) if operator == "IS NULL": return f"NOT EXISTS ({exists_sql} AND COALESCE(o.custodian, '') != '')", [] if operator == "IS NOT NULL": return f"EXISTS ({exists_sql} AND COALESCE(o.custodian, '') != '')", [] if operator in {"!=", "<>"}: assert isinstance(operand, dict) value = str(coerce_sql_literal("text", operand)) return f"NOT EXISTS ({exists_sql} AND COALESCE(o.custodian, '') = ?)", [value] clause, params = build_scalar_sql_filter_clause("o.custodian", field_def, operator, operand) return f"EXISTS ({exists_sql} AND {clause})", params def entity_display_label_sql(alias: str) -> str: return ( "CASE " f"WHEN COALESCE({alias}.display_name, '') != '' AND COALESCE({alias}.primary_email, '') != '' " f"THEN {alias}.display_name || ' <' || {alias}.primary_email || '>' " f"WHEN COALESCE({alias}.display_name, '') != '' THEN {alias}.display_name " f"WHEN COALESCE({alias}.primary_email, '') != '' THEN {alias}.primary_email " f"WHEN COALESCE({alias}.primary_phone, '') != '' THEN {alias}.primary_phone " f"ELSE 'Unknown Entity ' || {alias}.id END" ) def entity_role_exists_sql(alias: str, role: str, *, extra_clause: str | None = None) -> str: sql = ( "SELECT 1 " "FROM document_entities de " "JOIN entities e ON e.id = de.entity_id " f"WHERE de.document_id = {alias}.id " "AND e.canonical_status = 'active' " "AND de.role = ?" ) if extra_clause: sql += f" AND ({extra_clause})" return sql def entity_exact_values_for_filter(role: str, raw_value: object) -> list[str]: text = normalize_inline_whitespace(str(raw_value or "")) values: set[str] = set() if text: values.add(text.lower()) lookup_value = normalize_entity_lookup_text(text) if lookup_value: values.add(lookup_value.lower()) parsed_email = normalize_entity_email(text) if parsed_email: values.add(parsed_email.lower()) parsed_phone = normalize_entity_phone(text) if parsed_phone is not None: values.add(str(parsed_phone["display_value"]).lower()) values.add(str(parsed_phone["normalized_value"]).lower()) for candidate in parse_entity_candidates(text, role=role): for identifier in candidate.get("identifiers") or []: if not isinstance(identifier, dict): continue for key in ( "display_value", "normalized_value", "normalized_full_name", "normalized_sort_name", "identifier_name", ): value = normalize_inline_whitespace(str(identifier.get(key) or "")) if value: values.add(value.lower()) return sorted(values) def entity_exact_match_sql(role: str, operand: dict[str, object]) -> tuple[str, list[object]]: raw_value = coerce_sql_literal("text", operand) exact_values = entity_exact_values_for_filter(role, raw_value) if not exact_values: return "0", [] placeholders = ", ".join("?" for _ in exact_values) label_expr = entity_display_label_sql("e") expressions = [ f"LOWER(COALESCE({label_expr}, ''))", "LOWER(COALESCE(e.display_name, ''))", "LOWER(COALESCE(e.primary_email, ''))", "LOWER(COALESCE(e.primary_phone, ''))", "LOWER(COALESCE(ei.display_value, ''))", "LOWER(COALESCE(ei.normalized_value, ''))", "LOWER(COALESCE(ei.normalized_full_name, ''))", "LOWER(COALESCE(ei.normalized_sort_name, ''))", ] clause = " OR ".join(f"{expression} IN ({placeholders})" for expression in expressions) params: list[object] = [] for _ in expressions: params.extend(exact_values) return clause, params def entity_like_match_sql(operand: dict[str, object]) -> tuple[str, list[object]]: value = str(coerce_sql_literal("text", operand)) label_expr = entity_display_label_sql("e") expressions = [ f"LOWER(COALESCE({label_expr}, ''))", "LOWER(COALESCE(e.display_name, ''))", "LOWER(COALESCE(e.primary_email, ''))", "LOWER(COALESCE(e.primary_phone, ''))", "LOWER(COALESCE(ei.display_value, ''))", "LOWER(COALESCE(ei.normalized_value, ''))", "LOWER(COALESCE(ei.normalized_full_name, ''))", "LOWER(COALESCE(ei.normalized_sort_name, ''))", ] return " OR ".join(f"{expression} LIKE LOWER(?)" for expression in expressions), [value] * len(expressions) def build_entity_role_sql_filter_clause( alias: str, role: str, field_def: dict[str, object], operator: str, operand: object | None, ) -> tuple[str, list[object]]: supported = {"=", "!=", "<>", "IS NULL", "IS NOT NULL", "LIKE", "IN"} if operator not in supported: raise RetrieverError( f"Entity role field '{field_def['field_name']}' supports: =, !=, IS NULL, IS NOT NULL, LIKE, IN." ) exists_sql = entity_role_exists_sql(alias, role) if operator == "IS NULL": return f"NOT EXISTS ({exists_sql})", [role] if operator == "IS NOT NULL": return f"EXISTS ({exists_sql})", [role] if operator in {"=", "!=", "<>"}: assert isinstance(operand, dict) match_clause, match_params = entity_exact_match_sql(role, operand) matched_exists_sql = entity_role_exists_sql( alias, role, extra_clause=( "EXISTS (" "SELECT 1 FROM entity_identifiers ei " "WHERE ei.entity_id = e.id " f"AND ({match_clause})" ")" ), ) params = [role, *match_params] if operator == "=": return f"EXISTS ({matched_exists_sql})", params return f"NOT EXISTS ({matched_exists_sql})", params if operator == "LIKE": assert isinstance(operand, dict) match_clause, match_params = entity_like_match_sql(operand) matched_exists_sql = entity_role_exists_sql( alias, role, extra_clause=( "EXISTS (" "SELECT 1 FROM entity_identifiers ei " "WHERE ei.entity_id = e.id " f"AND ({match_clause})" ")" ), ) return f"EXISTS ({matched_exists_sql})", [role, *match_params] assert operator == "IN" assert isinstance(operand, list) if not operand: raise RetrieverError("IN requires at least one value.") if len(operand) > MAX_FILTER_IN_LIST_ITEMS: raise RetrieverError(f"IN (...) is capped at {MAX_FILTER_IN_LIST_ITEMS} values.") clauses: list[str] = [] params: list[object] = [] for literal in operand: match_clause, match_params = entity_exact_match_sql(role, literal) clauses.append(f"({match_clause})") params.extend(match_params) matched_exists_sql = entity_role_exists_sql( alias, role, extra_clause=( "EXISTS (" "SELECT 1 FROM entity_identifiers ei " "WHERE ei.entity_id = e.id " f"AND ({' OR '.join(clauses)})" ")" ), ) return f"EXISTS ({matched_exists_sql})", [role, *params] def build_entity_role_id_sql_filter_clause( alias: str, role: str, field_def: dict[str, object], operator: str, operand: object | None, ) -> tuple[str, list[object]]: supported = {"=", "!=", "<>", "IS NULL", "IS NOT NULL", "IN"} if operator not in supported: raise RetrieverError( f"Entity id field '{field_def['field_name']}' supports: =, !=, IS NULL, IS NOT NULL, IN." ) exists_sql = entity_role_exists_sql(alias, role) if operator == "IS NULL": return f"NOT EXISTS ({exists_sql})", [role] if operator == "IS NOT NULL": return f"EXISTS ({exists_sql})", [role] if operator == "IN": assert isinstance(operand, list) if len(operand) > MAX_FILTER_IN_LIST_ITEMS: raise RetrieverError(f"IN (...) is capped at {MAX_FILTER_IN_LIST_ITEMS} values.") values = [int(coerce_sql_literal("integer", literal)) for literal in operand] placeholders = ", ".join("?" for _ in values) matched_exists_sql = entity_role_exists_sql(alias, role, extra_clause=f"e.id IN ({placeholders})") return f"EXISTS ({matched_exists_sql})", [role, *values] assert isinstance(operand, dict) value = int(coerce_sql_literal("integer", operand)) matched_exists_sql = entity_role_exists_sql(alias, role, extra_clause="e.id = ?") params = [role, value] if operator == "=": return f"EXISTS ({matched_exists_sql})", params return f"NOT EXISTS ({matched_exists_sql})", params def build_raw_entity_metadata_sql_filter_clause( alias: str, field_def: dict[str, object], operator: str, operand: object | None, *, occurrence_column: str, occurrence_alias: str | None = None, ) -> tuple[str, list[object]]: raw_field_def = {**field_def, "field_type": "text"} if occurrence_alias is not None: return build_scalar_sql_filter_clause( f"{occurrence_alias}.{quote_identifier(occurrence_column)}", raw_field_def, operator, operand, ) ensure_sql_filter_operator_supported(raw_field_def, operator) exists_sql = ( "SELECT 1 " "FROM document_occurrences o " f"WHERE o.document_id = {alias}.id " "AND o.lifecycle_status = 'active'" ) column_expr = f"o.{quote_identifier(occurrence_column)}" if operator == "IS NULL": return f"NOT EXISTS ({exists_sql} AND COALESCE({column_expr}, '') != '')", [] if operator == "IS NOT NULL": return f"EXISTS ({exists_sql} AND COALESCE({column_expr}, '') != '')", [] if operator in {"!=", "<>"}: assert isinstance(operand, dict) value = str(coerce_sql_literal("text", operand)) return f"NOT EXISTS ({exists_sql} AND COALESCE({column_expr}, '') = ?)", [value] clause, params = build_scalar_sql_filter_clause(column_expr, raw_field_def, operator, operand) return f"EXISTS ({exists_sql} AND {clause})", params def virtual_field_sql_expression(alias: str, field_name: str) -> str: if field_name == "production_name": return ( "(SELECT p.production_name FROM productions p " f"WHERE p.id = {alias}.production_id)" ) if field_name == "is_attachment": return f"(CASE WHEN {attachment_child_filter_sql(alias)} THEN 1 ELSE 0 END)" if field_name == "has_attachments": return ( "(CASE WHEN EXISTS (" "SELECT 1 FROM documents child " f"WHERE child.parent_document_id = {alias}.id " f"AND {attachment_child_filter_sql('child')} " "AND child.lifecycle_status NOT IN ('missing', 'deleted')" ") THEN 1 ELSE 0 END)" ) raise RetrieverError(f"Unknown virtual filter: {field_name}") def build_sql_filter_clause( alias: str, field_def: dict[str, object], operator: str, operand: object | None, *, occurrence_alias: str | None = None, ) -> tuple[str, list[object]]: field_name = str(field_def["field_name"]) if occurrence_alias is not None and field_name in ENTITY_ROLE_FILTER_FIELDS: role = ENTITY_ROLE_FILTER_FIELDS[field_name] return build_raw_entity_metadata_sql_filter_clause( alias, field_def, operator, operand, occurrence_column=ROLE_OCCURRENCE_METADATA_COLUMNS[role], occurrence_alias=occurrence_alias, ) if field_name in ENTITY_ROLE_FILTER_FIELDS: return build_entity_role_sql_filter_clause( alias, ENTITY_ROLE_FILTER_FIELDS[field_name], field_def, operator, operand, ) if field_name in ENTITY_ROLE_ID_FILTER_FIELDS: return build_entity_role_id_sql_filter_clause( alias, ENTITY_ROLE_ID_FILTER_FIELDS[field_name], field_def, operator, operand, ) if field_name in RAW_ENTITY_METADATA_FILTER_FIELDS: return build_raw_entity_metadata_sql_filter_clause( alias, field_def, operator, operand, occurrence_column=RAW_ENTITY_METADATA_FILTER_FIELDS[field_name], occurrence_alias=occurrence_alias, ) if field_def.get("source") == "virtual": if field_name == "custodian": return build_custodian_sql_filter_clause( alias, field_def, operator, operand, occurrence_alias=occurrence_alias, ) if field_name == "dataset_name": return build_dataset_name_sql_filter_clause(alias, field_def, operator, operand) return build_scalar_sql_filter_clause(virtual_field_sql_expression(alias, field_name), field_def, operator, operand) if field_name in OCCURRENCE_FILTER_FIELDS: if occurrence_alias is not None: return build_scalar_sql_filter_clause( f"{occurrence_alias}.{quote_identifier(field_name)}", field_def, operator, operand, ) occurrence_clause, occurrence_params = build_scalar_sql_filter_clause( f"o.{quote_identifier(field_name)}", field_def, operator, operand, ) return ( "EXISTS (" "SELECT 1 FROM document_occurrences o " f"WHERE o.document_id = {alias}.id " "AND o.lifecycle_status = 'active' " f"AND {occurrence_clause}" ")", occurrence_params, ) target_alias = alias return build_scalar_sql_filter_clause( f"{target_alias}.{quote_identifier(field_name)}", field_def, operator, operand, ) def parse_sql_filter_literal(state: dict[str, object]) -> dict[str, object]: token = peek_filter_token(state) if token["kind"] != "literal": raise_filter_syntax_error(str(state["expression"]), int(token["start"]), "Expected a literal value") return consume_filter_token(state) def parse_sql_filter_predicate(state: dict[str, object]) -> tuple[str, list[object]]: identifier = expect_filter_token(state, "identifier", "Expected a field name") field_name = str(identifier["value"]) field_def = resolve_sql_filter_field(state["connection"], field_name) document_alias = str(state.get("document_alias", "d")) occurrence_alias = state.get("occurrence_alias") if match_filter_keyword(state, "IS"): operator = "IS NOT NULL" if match_filter_keyword(state, "NOT") else "IS NULL" token = expect_filter_token(state, "literal", "Expected NULL after IS / IS NOT") if token["literal_kind"] != "null": raise_filter_syntax_error(str(state["expression"]), int(token["start"]), "Expected NULL after IS / IS NOT") return build_sql_filter_clause( document_alias, field_def, operator, None, occurrence_alias=(str(occurrence_alias) if occurrence_alias is not None else None), ) negated_operator = match_filter_keyword(state, "NOT") if match_filter_keyword(state, "IN"): expect_filter_token(state, "lparen", "Expected '(' after IN") values: list[dict[str, object]] = [] while True: values.append(parse_sql_filter_literal(state)) if match_filter_token_kind(state, "comma") is None: break expect_filter_token(state, "rparen", "Expected ')' to close IN list") clause, params = build_sql_filter_clause( document_alias, field_def, "IN", values, occurrence_alias=(str(occurrence_alias) if occurrence_alias is not None else None), ) return (f"NOT ({clause})", params) if negated_operator else (clause, params) if match_filter_keyword(state, "BETWEEN"): left_value = parse_sql_filter_literal(state) if not match_filter_keyword(state, "AND"): token = peek_filter_token(state) raise_filter_syntax_error(str(state["expression"]), int(token["start"]), "Expected AND in BETWEEN expression") right_value = parse_sql_filter_literal(state) clause, params = build_sql_filter_clause( document_alias, field_def, "BETWEEN", (left_value, right_value), occurrence_alias=(str(occurrence_alias) if occurrence_alias is not None else None), ) return (f"NOT ({clause})", params) if negated_operator else (clause, params) if match_filter_keyword(state, "LIKE"): value = parse_sql_filter_literal(state) clause, params = build_sql_filter_clause( document_alias, field_def, "LIKE", value, occurrence_alias=(str(occurrence_alias) if occurrence_alias is not None else None), ) return (f"NOT ({clause})", params) if negated_operator else (clause, params) if negated_operator: token = peek_filter_token(state) raise_filter_syntax_error(str(state["expression"]), int(token["start"]), "Expected IN, BETWEEN, or LIKE after NOT") operator_token = expect_filter_token(state, "operator", "Expected an operator after the field name") value = parse_sql_filter_literal(state) return build_sql_filter_clause( document_alias, field_def, str(operator_token["value"]).upper(), value, occurrence_alias=(str(occurrence_alias) if occurrence_alias is not None else None), ) def parse_sql_filter_primary(state: dict[str, object]) -> tuple[str, list[object]]: if match_filter_token_kind(state, "lparen") is not None: clause, params = parse_sql_filter_or_expression(state) expect_filter_token(state, "rparen", "Expected ')' to close grouped expression") return clause, params return parse_sql_filter_predicate(state) def parse_sql_filter_not_expression(state: dict[str, object]) -> tuple[str, list[object]]: if match_filter_keyword(state, "NOT"): clause, params = parse_sql_filter_not_expression(state) return f"NOT ({clause})", params return parse_sql_filter_primary(state) def parse_sql_filter_and_expression(state: dict[str, object]) -> tuple[str, list[object]]: clause, params = parse_sql_filter_not_expression(state) while match_filter_keyword(state, "AND"): right_clause, right_params = parse_sql_filter_not_expression(state) clause = f"({clause}) AND ({right_clause})" params.extend(right_params) return clause, params def parse_sql_filter_or_expression(state: dict[str, object]) -> tuple[str, list[object]]: clause, params = parse_sql_filter_and_expression(state) while match_filter_keyword(state, "OR"): right_clause, right_params = parse_sql_filter_and_expression(state) clause = f"({clause}) OR ({right_clause})" params.extend(right_params) return clause, params def compile_sql_filter_expression( connection: sqlite3.Connection, expression: str, *, document_alias: str = "d", occurrence_alias: str | None = None, ) -> tuple[str, list[object]]: normalized_expression = expression.strip() if not normalized_expression: raise RetrieverError("Filter expression cannot be empty.") state: dict[str, object] = { "connection": connection, "expression": normalized_expression, "tokens": tokenize_sql_filter_expression(normalized_expression), "index": 0, "document_alias": document_alias, "occurrence_alias": occurrence_alias, } clause, params = parse_sql_filter_or_expression(state) trailing_token = peek_filter_token(state) if trailing_token["kind"] != "eof": raise_filter_syntax_error( normalized_expression, int(trailing_token["start"]), "Unexpected token after complete expression", ) return clause, params def normalize_sql_filter_expressions(raw_filters: object | None) -> list[str]: if raw_filters is None: return [] if isinstance(raw_filters, str): return [raw_filters] if raw_filters.strip() else [] if not isinstance(raw_filters, list): raise RetrieverError("Filters must be provided as strings or repeatable --filter arguments.") expressions: list[str] = [] for item in raw_filters: if isinstance(item, str): if item.strip(): expressions.append(item) continue if isinstance(item, (list, tuple)): expression = " ".join(str(part) for part in item if normalize_inline_whitespace(str(part or ""))) if expression: expressions.append(expression) continue raise RetrieverError("Unsupported filter payload shape.") return expressions def build_sql_like_search_filters( connection: sqlite3.Connection, raw_filters: object | None, ) -> tuple[list[str], list[str], list[object]]: expressions = normalize_sql_filter_expressions(raw_filters) clauses = base_document_search_clauses() params: list[object] = [] for expression in expressions: clause, clause_params = compile_sql_filter_expression(connection, expression) clauses.append(f"({clause})") params.extend(clause_params) return expressions, clauses, params def metadata_snippet(row: sqlite3.Row) -> str: parts = [ row["control_number"], row["begin_bates"], row["end_bates"], row["content_type"], document_custodian_display_text_from_row(row), row["source_rel_path"], row["source_folder_path"], row["title"], row["subject"], row["author"], row["participants"], row["recipients"], ] text = normalize_whitespace(" | ".join(part for part in parts if part)) return text[:220] def query_terms(query: str) -> list[str]: return list(dict.fromkeys(re.findall(r"[A-Za-z0-9_]+", query.lower()))) def make_snippet(text: str, query: str) -> str: normalized = normalize_whitespace(text) if not normalized: return "" lower_text = normalized.lower() start = 0 for term in query_terms(query): index = lower_text.find(term) if index != -1: start = max(0, index - 80) break end = min(len(normalized), start + 220) snippet = normalized[start:end] if start > 0: snippet = "..." + snippet if end < len(normalized): snippet = snippet + "..." return snippet def fetch_documents_by_ids(connection: sqlite3.Connection, document_ids: list[int]) -> dict[int, sqlite3.Row]: if not document_ids: return {} placeholders = ", ".join("?" for _ in document_ids) rows = connection.execute( f"SELECT * FROM documents WHERE id IN ({placeholders})", document_ids, ).fetchall() return {int(row["id"]): row for row in rows} def build_occurrence_scope_filters( connection: sqlite3.Connection, raw_filters: object | None, ) -> tuple[list[str], list[object]]: clauses = ["o.lifecycle_status = 'active'"] params: list[object] = [] for expression in normalize_sql_filter_expressions(raw_filters): clause, clause_params = compile_sql_filter_expression( connection, expression, document_alias="d", occurrence_alias="o", ) clauses.append(f"({clause})") params.extend(clause_params) return clauses, params def preferred_occurrence_for_document( connection: sqlite3.Connection, document_id: int, occurrence_scope_clauses: list[str], occurrence_scope_params: list[object], ) -> sqlite3.Row | None: scoped_rows = connection.execute( f""" SELECT o.* FROM document_occurrences o JOIN documents d ON d.id = o.document_id WHERE o.document_id = ? AND {' AND '.join(occurrence_scope_clauses)} ORDER BY o.id ASC """, [document_id, *occurrence_scope_params], ).fetchall() preferred = select_preferred_occurrence(scoped_rows) if preferred is not None: return preferred return select_preferred_occurrence(active_occurrence_rows_for_document(connection, document_id)) def preferred_occurrences_by_document( connection: sqlite3.Connection, document_ids: list[int], occurrence_scope_clauses: list[str], occurrence_scope_params: list[object], ) -> dict[int, sqlite3.Row]: preferred: dict[int, sqlite3.Row] = {} for document_id in document_ids: occurrence_row = preferred_occurrence_for_document( connection, document_id, occurrence_scope_clauses, occurrence_scope_params, ) if occurrence_row is not None: preferred[document_id] = occurrence_row return preferred def document_path_payload( paths: dict[str, Path], connection: sqlite3.Connection, row: sqlite3.Row, *, occurrence_row: sqlite3.Row | None = None, include_preview_targets: bool = True, prefer_preview_label: str | None = "message", ) -> dict[str, object]: preview_target = default_preview_target(paths, row, connection) effective_rel_path = str(occurrence_row["rel_path"]) if occurrence_row is not None else str(row["rel_path"]) preview_targets: list[dict[str, object]] | None = None normalized_preferred_label = normalize_whitespace(str(prefer_preview_label or "")).lower() if include_preview_targets or normalized_preferred_label: preview_targets = collect_preview_targets(paths, int(row["id"]), effective_rel_path, connection) if normalized_preferred_label and preview_targets: preferred_target = next( ( target for target in preview_targets if normalize_whitespace(str(target.get("label") or "")).lower() == normalized_preferred_label ), None, ) if preferred_target is not None: preview_target = preferred_target if ( occurrence_row is not None and str(preview_target.get("preview_type") or "") == "native" and str(preview_target.get("rel_path") or "") == str(row["rel_path"]) ): effective_abs_path = str(document_absolute_path(paths, effective_rel_path)) preview_target = build_preview_target_payload( rel_path=effective_rel_path, abs_path=effective_abs_path, preview_type="native", label=None, ordinal=0, ) payload = { "rel_path": effective_rel_path, "abs_path": str(document_absolute_path(paths, effective_rel_path)), "preview_rel_path": preview_target["rel_path"], "preview_abs_path": preview_target["abs_path"], "preview_file_rel_path": preview_target["file_rel_path"], "preview_file_abs_path": preview_target["file_abs_path"], "preview_target_fragment": preview_target["target_fragment"], } if include_preview_targets: payload["preview_targets"] = preview_targets or collect_preview_targets(paths, int(row["id"]), effective_rel_path, connection) return payload def resolve_conversation_preview_rel_path(paths: dict[str, Path], conversation_id: int) -> str: preferred_rel_path = conversation_preview_full_rel_path(conversation_id) preferred_abs_path = paths["state_dir"] / preferred_rel_path if preferred_abs_path.exists(): return preferred_rel_path preview_dir = paths["state_dir"] / conversation_preview_base_path(conversation_id) if preview_dir.exists(): segment_paths = sorted( path.relative_to(paths["state_dir"]).as_posix() for path in preview_dir.glob("segment-*.html") if path.is_file() ) if segment_paths: return segment_paths[0] toc_rel_path = conversation_preview_toc_rel_path(conversation_id) if (paths["state_dir"] / toc_rel_path).exists(): return toc_rel_path return preferred_rel_path def single_document_conversation_preview_target( paths: dict[str, Path], connection: sqlite3.Connection, conversation_id: int, ) -> dict[str, object] | None: documents = load_preview_documents(connection, paths, conversation_id=conversation_id) if len(documents) != 1: return None document_row = connection.execute( "SELECT * FROM documents WHERE id = ?", (int(documents[0]["id"]),), ).fetchone() if document_row is None: return None preview_target = dict(default_preview_target(paths, document_row, connection)) preview_target["label"] = "conversation" preview_target["ordinal"] = 0 return preview_target def conversation_path_payload( paths: dict[str, Path], connection: sqlite3.Connection, conversation_id: int, *, include_preview_targets: bool = True, ) -> dict[str, object]: rel_preview_path = resolve_conversation_preview_rel_path(paths, conversation_id) abs_preview_path = paths["state_dir"] / rel_preview_path if abs_preview_path.exists(): preview_target = build_preview_target_payload( rel_path=str(Path(INTERNAL_REL_PATH_PREFIX) / rel_preview_path), abs_path=str(abs_preview_path), preview_type="html", label="conversation", ordinal=0, ) else: preview_target = single_document_conversation_preview_target( paths, connection, conversation_id, ) or build_preview_target_payload( rel_path=str(Path(INTERNAL_REL_PATH_PREFIX) / rel_preview_path), abs_path=str(abs_preview_path), preview_type="html", label="conversation", ordinal=0, ) payload = { "rel_path": preview_target["rel_path"], "abs_path": preview_target["abs_path"], "preview_rel_path": preview_target["rel_path"], "preview_abs_path": preview_target["abs_path"], "preview_file_rel_path": preview_target["file_rel_path"], "preview_file_abs_path": preview_target["file_abs_path"], "preview_target_fragment": preview_target["target_fragment"], } if include_preview_targets: payload["preview_targets"] = [preview_target] return payload def fetch_attachment_counts( connection: sqlite3.Connection, parent_ids: list[int], ) -> dict[int, int]: if not parent_ids: return {} placeholders = ", ".join("?" for _ in parent_ids) rows = connection.execute( f""" SELECT parent_document_id, COUNT(*) AS attachment_count FROM documents WHERE parent_document_id IN ({placeholders}) AND COALESCE(child_document_kind, ?) = ? AND lifecycle_status NOT IN ('missing', 'deleted') GROUP BY parent_document_id """, [*parent_ids, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT], ).fetchall() return {int(row["parent_document_id"]): int(row["attachment_count"]) for row in rows} def fetch_attachment_summaries( connection: sqlite3.Connection, paths: dict[str, Path], parent_ids: list[int], ) -> dict[int, list[dict[str, object]]]: if not parent_ids: return {} placeholders = ", ".join("?" for _ in parent_ids) rows = connection.execute( f""" SELECT * FROM documents WHERE parent_document_id IN ({placeholders}) AND COALESCE(child_document_kind, ?) = ? AND lifecycle_status NOT IN ('missing', 'deleted') ORDER BY parent_document_id ASC, id ASC """, [*parent_ids, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT], ).fetchall() grouped: dict[int, list[dict[str, object]]] = defaultdict(list) for row in rows: grouped[int(row["parent_document_id"])].append( { "id": int(row["id"]), "control_number": row["control_number"], "file_name": row["file_name"], "file_type": row["file_type"], "content_type": row["content_type"], **document_path_payload(paths, connection, row), "source_kind": row["source_kind"], "begin_bates": row["begin_bates"], "control_number_attachment_sequence": row["control_number_attachment_sequence"], } ) for parent_id, items in grouped.items(): grouped[parent_id] = sorted( items, key=lambda item: ( 0 if item.get("source_kind") == PRODUCTION_SOURCE_KIND else 1, bates_sort_key(item.get("begin_bates") or item.get("control_number")) if item.get("source_kind") == PRODUCTION_SOURCE_KIND else (0, "", int(item.get("control_number_attachment_sequence") or 0), ""), item["id"], ), ) for item in grouped[parent_id]: item.pop("source_kind", None) item.pop("begin_bates", None) item.pop("control_number_attachment_sequence", None) return grouped def fetch_child_document_summaries( connection: sqlite3.Connection, paths: dict[str, Path], parent_ids: list[int], ) -> dict[int, list[dict[str, object]]]: if not parent_ids: return {} placeholders = ", ".join("?" for _ in parent_ids) rows = connection.execute( f""" SELECT * FROM documents WHERE parent_document_id IN ({placeholders}) AND COALESCE(child_document_kind, ?) != ? AND lifecycle_status NOT IN ('missing', 'deleted') ORDER BY parent_document_id ASC, COALESCE(date_created, '') ASC, id ASC """, [*parent_ids, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT], ).fetchall() grouped: dict[int, list[dict[str, object]]] = defaultdict(list) for row in rows: grouped[int(row["parent_document_id"])].append( { "id": int(row["id"]), "control_number": row["control_number"], "file_name": row["file_name"], "file_type": row["file_type"], "content_type": row["content_type"], **document_path_payload(paths, connection, row), "child_document_kind": row["child_document_kind"], "title": row["title"], "date_created": row["date_created"], "date_modified": row["date_modified"], "root_message_key": row["root_message_key"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], } ) return grouped def fetch_child_document_counts( connection: sqlite3.Connection, parent_ids: list[int], ) -> dict[int, int]: if not parent_ids: return {} placeholders = ", ".join("?" for _ in parent_ids) rows = connection.execute( f""" SELECT parent_document_id, COUNT(*) AS child_document_count FROM documents WHERE parent_document_id IN ({placeholders}) AND COALESCE(child_document_kind, ?) != ? AND lifecycle_status NOT IN ('missing', 'deleted') GROUP BY parent_document_id """, [*parent_ids, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT], ).fetchall() return {int(row["parent_document_id"]): int(row["child_document_count"]) for row in rows} def fetch_parent_summaries( connection: sqlite3.Connection, child_rows: list[sqlite3.Row], ) -> dict[int, dict[str, object]]: parent_ids = sorted( { int(row["parent_document_id"]) for row in child_rows if row["parent_document_id"] is not None } ) if not parent_ids: return {} placeholders = ", ".join("?" for _ in parent_ids) rows = connection.execute( f""" SELECT id, control_number, rel_path, file_name, subject, author, date_created FROM documents WHERE id IN ({placeholders}) """, parent_ids, ).fetchall() return { int(row["id"]): { "id": int(row["id"]), "control_number": row["control_number"], "rel_path": row["rel_path"], "file_name": row["file_name"], "subject": row["subject"], "author": row["author"], "date_created": row["date_created"], } for row in rows } def fetch_production_names(connection: sqlite3.Connection, rows: list[sqlite3.Row]) -> dict[int, str]: production_ids = sorted({int(row["production_id"]) for row in rows if row["production_id"] is not None}) if not production_ids: return {} placeholders = ", ".join("?" for _ in production_ids) result_rows = connection.execute( f"SELECT id, production_name FROM productions WHERE id IN ({placeholders})", production_ids, ).fetchall() return {int(row["id"]): str(row["production_name"]) for row in result_rows} def fetch_dataset_memberships_for_document_ids( connection: sqlite3.Connection, document_ids: list[int], ) -> dict[int, dict[str, list[object]]]: normalized_document_ids = sorted({int(document_id) for document_id in document_ids}) if not normalized_document_ids: return {} placeholders = ", ".join("?" for _ in normalized_document_ids) result_rows = connection.execute( f""" SELECT dd.document_id, ds.id AS dataset_id, ds.dataset_name FROM dataset_documents dd JOIN datasets ds ON ds.id = dd.dataset_id WHERE dd.document_id IN ({placeholders}) ORDER BY dd.document_id ASC, LOWER(ds.dataset_name) ASC, ds.id ASC """, normalized_document_ids, ).fetchall() memberships: dict[int, dict[str, list[object]]] = defaultdict(lambda: {"ids": [], "names": []}) for row in result_rows: payload = memberships[int(row["document_id"])] dataset_id = int(row["dataset_id"]) dataset_name = str(row["dataset_name"]) if dataset_id not in payload["ids"]: payload["ids"].append(dataset_id) if dataset_name not in payload["names"]: payload["names"].append(dataset_name) return memberships def fetch_document_dataset_memberships( connection: sqlite3.Connection, rows: list[sqlite3.Row], ) -> dict[int, dict[str, list[object]]]: return fetch_dataset_memberships_for_document_ids( connection, [int(row["id"]) for row in rows], ) def load_conversation_summary_documents( connection: sqlite3.Connection, conversation_ids: list[int], ) -> dict[int, list[dict[str, object]]]: normalized_conversation_ids = sorted({int(conversation_id) for conversation_id in conversation_ids}) if not normalized_conversation_ids: return {} placeholders = ", ".join("?" for _ in normalized_conversation_ids) rows = connection.execute( f""" SELECT id, conversation_id, parent_document_id, author, participants, recipients, date_created, date_modified, source_rel_path, rel_path FROM documents WHERE conversation_id IN ({placeholders}) AND lifecycle_status NOT IN ('missing', 'deleted') AND COALESCE(child_document_kind, '') != ? ORDER BY conversation_id ASC, id ASC """, [*normalized_conversation_ids, CHILD_DOCUMENT_KIND_ATTACHMENT], ).fetchall() grouped: dict[int, list[dict[str, object]]] = defaultdict(list) for row in rows: grouped[int(row["conversation_id"])].append( { "id": int(row["id"]), "conversation_id": int(row["conversation_id"]), "parent_document_id": row["parent_document_id"], "author": row["author"], "participants": row["participants"], "recipients": row["recipients"], "date_created": row["date_created"], "date_modified": row["date_modified"], "source_rel_path": row["source_rel_path"], "rel_path": row["rel_path"], } ) return dict(grouped) def chunk_preview_text(text: object, *, max_chars: int = 220) -> str: normalized = normalize_whitespace(str(text or "")) if len(normalized) <= max_chars: return normalized return normalized[: max_chars - 3].rstrip() + "..." COMPACT_METADATA_FIELDS = ( "author", "content_type", "custodian", "date_created", "date_modified", "participants", "recipients", "subject", "title", "updated_at", ) def payload_has_meaningful_value(value: object) -> bool: return value not in (None, "", [], {}) def compact_metadata_payload(metadata: object) -> dict[str, object]: if not isinstance(metadata, dict): return {} return { key: metadata[key] for key in COMPACT_METADATA_FIELDS if key in metadata and payload_has_meaningful_value(metadata[key]) } def compact_search_result_payload(item: dict[str, object]) -> dict[str, object]: compact: dict[str, object] = {"id": item["id"]} for key in ("control_number", "file_name", "file_type", "preview_abs_path", "snippet", "rank"): if key in item and payload_has_meaningful_value(item[key]): compact[key] = item[key] for key in ( "conversation_id", "conversation_type", "title", "participants", "first_activity", "last_activity", "document_count", "matching_document_count", "source_kind", ): if key in item and payload_has_meaningful_value(item[key]): compact[key] = item[key] metadata = compact_metadata_payload(item.get("metadata")) if metadata: compact["metadata"] = metadata if payload_has_meaningful_value(item.get("dataset_name")): compact["dataset_name"] = item["dataset_name"] elif payload_has_meaningful_value(item.get("dataset_names")): compact["dataset_names"] = item["dataset_names"] if payload_has_meaningful_value(item.get("production_name")): compact["production_name"] = item["production_name"] if int(item.get("attachment_count") or 0) > 0: compact["attachment_count"] = int(item["attachment_count"]) if int(item.get("child_document_count") or 0) > 0: compact["child_document_count"] = int(item["child_document_count"]) if payload_has_meaningful_value(item.get("parent")): compact["parent"] = item["parent"] if "display_values" in item: compact["display_values"] = item["display_values"] return compact def compact_document_overview_payload(item: object) -> object: if not isinstance(item, dict): return item compact: dict[str, object] = {"document_id": item["document_id"]} for key in ("control_number", "file_name", "file_type", "preview_rel_path", "preview_abs_path"): if key in item and payload_has_meaningful_value(item[key]): compact[key] = item[key] metadata = compact_metadata_payload(item.get("metadata")) if metadata: compact["metadata"] = metadata if payload_has_meaningful_value(item.get("dataset_name")): compact["dataset_name"] = item["dataset_name"] elif payload_has_meaningful_value(item.get("dataset_names")): compact["dataset_names"] = item["dataset_names"] if payload_has_meaningful_value(item.get("production_name")): compact["production_name"] = item["production_name"] if int(item.get("attachment_count") or 0) > 0: compact["attachment_count"] = int(item["attachment_count"]) if int(item.get("child_document_count") or 0) > 0: compact["child_document_count"] = int(item["child_document_count"]) if payload_has_meaningful_value(item.get("parent")): compact["parent"] = item["parent"] return compact def compact_search_payload(payload: dict[str, object]) -> dict[str, object]: compact_payload = { "query": payload["query"], "filters": payload["filters"], "sort": payload["sort"], "order": payload["order"], "browse_mode": normalize_browse_mode(payload.get("browse_mode")), "page": payload["page"], "per_page": payload["per_page"], "total_hits": payload["total_hits"], "total_pages": payload["total_pages"], "results": [compact_search_result_payload(item) for item in payload["results"]], } if payload_has_meaningful_value(payload.get("scope")): compact_payload["scope"] = payload["scope"] if payload_has_meaningful_value(payload.get("header")): compact_payload["header"] = payload["header"] if payload_has_meaningful_value(payload.get("display")): compact_payload["display"] = payload["display"] if payload_has_meaningful_value(payload.get("warnings")): compact_payload["warnings"] = payload["warnings"] if payload_has_meaningful_value(payload.get("rendered_markdown")): compact_payload["rendered_markdown"] = payload["rendered_markdown"] return compact_payload def compact_get_doc_payload(payload: dict[str, object]) -> dict[str, object]: compact = { "status": payload["status"], "document": compact_document_overview_payload(payload["document"]), "chunk_count": payload["chunk_count"], "include_text": payload["include_text"], "chunks": payload["chunks"], } if payload.get("text_summary") is not None: compact["text_summary"] = payload["text_summary"] return compact def compact_search_chunk_result_payload(item: dict[str, object]) -> dict[str, object]: compact: dict[str, object] = { "document_id": item["document_id"], "chunk_index": item["chunk_index"], "char_start": item["char_start"], "char_end": item["char_end"], "citation": item["citation"], } for key in ("control_number", "file_name", "file_type", "token_estimate", "snippet", "rank"): if key in item and payload_has_meaningful_value(item[key]): compact[key] = item[key] metadata = compact_metadata_payload(item.get("metadata")) if metadata: compact["metadata"] = metadata if payload_has_meaningful_value(item.get("dataset_name")): compact["dataset_name"] = item["dataset_name"] elif payload_has_meaningful_value(item.get("dataset_names")): compact["dataset_names"] = item["dataset_names"] if payload_has_meaningful_value(item.get("production_name")): compact["production_name"] = item["production_name"] if payload_has_meaningful_value(item.get("parent")): compact["parent"] = item["parent"] return compact def compact_search_chunks_payload(payload: dict[str, object]) -> dict[str, object]: if "results" not in payload: return payload compact_payload = { "query": payload["query"], "filters": payload["filters"], "sort": payload["sort"], "order": payload["order"], "top_k": payload["top_k"], "per_doc_cap": payload["per_doc_cap"], "total_matches": payload["total_matches"], "results": [compact_search_chunk_result_payload(item) for item in payload["results"]], } if payload_has_meaningful_value(payload.get("scope")): compact_payload["scope"] = payload["scope"] if payload.get("selected_from_scope"): compact_payload["selected_from_scope"] = True return compact_payload def prepare_cli_payload(command: str, payload: dict[str, object], *, verbose: bool = False) -> dict[str, object]: if verbose: return payload if command == "search": return compact_search_payload(payload) if command == "get-doc": return compact_get_doc_payload(payload) if command == "search-chunks": return compact_search_chunks_payload(payload) return payload CLI_OUTPUT_MODE = "json" LONG_TASK_PROGRESS_INTERVAL_SECONDS = 60.0 def set_cli_output_mode(mode: str) -> None: global CLI_OUTPUT_MODE CLI_OUTPUT_MODE = "human" if normalize_inline_whitespace(mode).lower() == "human" else "json" def cli_run_payload(payload: dict[str, object]) -> dict[str, object]: run = payload.get("run") if isinstance(run, dict): return run return payload def cli_work_item_counts(payload: dict[str, object]) -> dict[str, object]: counts = payload.get("counts") if not isinstance(counts, dict): counts = cli_run_payload(payload).get("counts") if not isinstance(counts, dict): return {} work_items = counts.get("work_items") return dict(work_items) if isinstance(work_items, dict) else {} def cli_counts_by_unit_type(payload: dict[str, object]) -> dict[str, dict[str, object]]: counts = payload.get("counts") if not isinstance(counts, dict): counts = cli_run_payload(payload).get("counts") if not isinstance(counts, dict): return {} by_unit_type = counts.get("by_unit_type") if not isinstance(by_unit_type, dict): return {} normalized: dict[str, dict[str, object]] = {} for unit_type, unit_counts in by_unit_type.items(): if isinstance(unit_counts, dict): normalized[str(unit_type)] = dict(unit_counts) return normalized def human_count_label(value: object, singular: str, plural: str | None = None) -> str: count = int(value or 0) noun = singular if count == 1 else (plural or singular + "s") return f"{count} {noun}" def humanize_unit_type(unit_type: str) -> str: return unit_type.replace("_", " ") def first_status_detail(status_report: dict[str, object]) -> str | None: for value in status_report.values(): if not isinstance(value, dict): continue if str(value.get("status") or "").lower() != "fail": continue detail = normalize_inline_whitespace(str(value.get("detail") or value.get("error") or "")) if detail: return detail return None def render_workspace_human_output(payload: dict[str, object]) -> str: action = normalize_inline_whitespace(str(payload.get("action") or "status")) or "status" status_report = payload.get("status_report") if isinstance(payload.get("status_report"), dict) else payload workspace_root = normalize_inline_whitespace( str( payload.get("workspace_root") or status_report.get("workspace_root") or (status_report.get("workspace") or {}).get("root") or "" ) ) overall = normalize_inline_whitespace(str(status_report.get("overall") or payload.get("status") or "")) tool_version = normalize_inline_whitespace( str( status_report.get("tool_version") or (payload.get("initialization") or {}).get("tool_version") or (payload.get("tool_update") or {}).get("new_tool_version") or TOOL_VERSION ) ) schema_version = normalize_inline_whitespace( str( status_report.get("schema_version") or status_report.get("workspace_schema_version") or (payload.get("initialization") or {}).get("schema_version") or SCHEMA_VERSION ) ) runtime_status = normalize_inline_whitespace( str( (payload.get("runtime_init") or {}).get("status") or (status_report.get("plugin_runtime") or {}).get("status") or "" ) ) or "unknown" workspace_inventory = ( status_report.get("workspace_inventory") if isinstance(status_report.get("workspace_inventory"), dict) else {} ) documents_total = int(workspace_inventory.get("documents_total") or 0) ready = overall.lower() in {"pass", "ready", "ok"} action_phrase = { "init": "Done. Workspace initialized and ready", "status": "Workspace ready", "update": "Done. Workspace updated and ready", "doctor": "Workspace doctor completed", }.get(action, "Workspace ready") if ready: return ( f"{action_phrase} at {workspace_root} " f"(tool v{tool_version}, schema {schema_version}, runtime {runtime_status}, " f"{human_count_label(documents_total, 'document')})." ) detail = first_status_detail(status_report) or normalize_inline_whitespace(str(payload.get("status") or overall)) root_suffix = f" at {workspace_root}" if workspace_root else "" return f"Workspace {action} failed{root_suffix}: {detail or 'unknown error'}." def render_ingest_human_output(payload: dict[str, object]) -> str: run = cli_run_payload(payload) counts = cli_work_item_counts(payload) by_unit_type = cli_counts_by_unit_type(payload) status = normalize_inline_whitespace(str(payload.get("status") or run.get("status") or "")) phase = normalize_inline_whitespace(str(payload.get("phase") or run.get("phase") or status)) run_id = normalize_inline_whitespace(str(payload.get("run_id") or run.get("run_id") or "")) committed = int(counts.get("committed") or 0) pending = int(counts.get("pending") or 0) prepared = int(counts.get("prepared") or 0) leased = int(counts.get("leased") or 0) failed = int(counts.get("failed") or 0) canceled = int(counts.get("cancelled") or 0) more_work_remaining = bool(payload.get("more_work_remaining")) line: str if status == "completed": line = f"Done. Ingest completed — {committed} committed, {failed} failed" if canceled: line += f", {canceled} cancelled" elif status == "canceled": line = f"Ingest canceled — {committed} committed, {failed} failed" elif status == "none": line = "No ingest run found." else: line = ( f"Ingest {phase or 'in progress'} — {committed} committed, {pending} pending, " f"{prepared} prepared, {leased} leased, {failed} failed" ) if run_id: line += f" (run_id {run_id})" lines = [line + "."] unit_notes: list[str] = [] for unit_type, unit_counts in sorted(by_unit_type.items()): unit_committed = int(unit_counts.get("committed") or 0) unit_pending = int(unit_counts.get("pending") or 0) unit_failed = int(unit_counts.get("failed") or 0) if unit_committed == 0 and unit_pending == 0 and unit_failed == 0: continue fragments = [] if unit_committed: fragments.append(f"{unit_committed} committed") if unit_pending: fragments.append(f"{unit_pending} pending") if unit_failed: fragments.append(f"{unit_failed} failed") if fragments: unit_notes.append(f"{humanize_unit_type(unit_type)} {', '.join(fragments)}") if unit_notes: lines.append("By type: " + "; ".join(unit_notes[:4]) + ("." if len(unit_notes) <= 4 else "; ...")) if more_work_remaining: next_commands = payload.get("next_recommended_commands") if isinstance(next_commands, list) and next_commands: lines.append(f"Next: {next_commands[0]}") return "\n".join(lines) def render_export_human_output(payload: dict[str, object]) -> str: if str(payload.get("slash_command") or "") == "/export status" and isinstance(payload.get("exports"), dict): active_exports = payload.get("active_exports") if isinstance(payload.get("active_exports"), list) else [] if not active_exports: run_id = normalize_inline_whitespace(str(payload.get("run_id") or "")) if run_id: return f"No export run found for run_id {run_id}." return "No active exports." lines = ["Active exports:"] for item in active_exports: if not isinstance(item, dict): continue label = normalize_inline_whitespace(str(item.get("export_label") or "export")) status = normalize_inline_whitespace(str(item.get("status") or item.get("phase") or "active")) run_id = normalize_inline_whitespace(str(item.get("run_id") or "")) suffix = f" (run_id {run_id})" if run_id else "" lines.append(f"- {label}: {status}{suffix}") return "\n".join(lines) run = cli_run_payload(payload) counts = cli_work_item_counts(payload) export_kind = normalize_inline_whitespace( str(payload.get("export_label") or run.get("export_kind") or payload.get("export_kind") or "export") ) if export_kind == "csv": export_kind = "table" status = normalize_inline_whitespace(str(payload.get("status") or run.get("status") or "")) phase = normalize_inline_whitespace(str(payload.get("phase") or run.get("phase") or status)) run_id = normalize_inline_whitespace(str(payload.get("run_id") or run.get("run_id") or "")) completed = int(counts.get("completed") or 0) pending = int(counts.get("pending") or 0) running = int(counts.get("running") or 0) failed = int(counts.get("failed") or 0) output_rel_path = normalize_inline_whitespace(str(run.get("output_rel_path") or "")) output_path = normalize_inline_whitespace(str(run.get("output_path") or "")) more_work_remaining = bool(payload.get("more_work_remaining")) or status not in EXPORT_RUN_TERMINAL_STATUSES if status == "completed": line = f"Done. {export_kind.title()} export completed — {completed} completed, {failed} failed" elif status == "failed": line = f"{export_kind.title()} export failed — {completed} completed, {failed} failed" elif status == "none": return f"No {export_kind} export run found." else: line = ( f"{export_kind.title()} export {phase or 'in progress'} — {completed} completed, " f"{pending} pending, {running} running, {failed} failed" ) if run_id: line += f" (run_id {run_id})" lines = [line + "."] if output_rel_path or output_path: lines.append(f"Output: {output_rel_path or output_path}") if more_work_remaining: next_commands = payload.get("next_recommended_commands") or run.get("next_recommended_commands") if isinstance(next_commands, list) and next_commands: lines.append(f"Next: {next_commands[0]}") return "\n".join(lines) def render_get_doc_human_output(payload: dict[str, object]) -> str: document = payload.get("document") if isinstance(payload.get("document"), dict) else {} document_id = int(document.get("document_id") or 0) title = normalize_inline_whitespace( str(document.get("control_number") or document.get("file_name") or f"Document {document_id}") ) preview_abs_path = normalize_inline_whitespace(str(document.get("preview_abs_path") or "")) preview_target = preview_abs_path if preview_target and not preview_target.startswith("file://"): preview_target = "file://" + preview_target lines = [f"Document {document_id}: {title}."] if preview_target: lines.append(f"Preview: {preview_target}") text_summary = normalize_inline_whitespace(str(payload.get("text_summary") or "")) if text_summary: lines.append(f"Text: {text_summary}") return "\n".join(lines) def render_cli_human_output(command: str, payload: dict[str, object]) -> str | None: if command == "workspace": return render_workspace_human_output(payload) if command in { "ingest", "ingest-start", "ingest-status", "ingest-cancel", "ingest-run-step", "ingest-plan-step", "ingest-prepare-step", "ingest-commit-step", "ingest-finalize-step", }: return render_ingest_human_output(payload) if command in {"export-csv-start", "export-csv-status", "export-csv-run-step", "export-archive-start", "export-archive-status", "export-archive-run-step"}: return render_export_human_output(payload) if command == "slash" and str(payload.get("slash_command") or "").startswith("/export"): return render_export_human_output(payload) if command == "get-doc": return render_get_doc_human_output(payload) return None def render_progress_line(command: str, payload: dict[str, object]) -> str | None: rendered = render_cli_human_output(command, payload) if isinstance(rendered, str): lines = [normalize_inline_whitespace(line) for line in rendered.splitlines()] first_line = next((line for line in lines if line), "") if first_line: return first_line run = cli_run_payload(payload) status = normalize_inline_whitespace(str(payload.get("status") or run.get("status") or "")) phase = normalize_inline_whitespace(str(payload.get("phase") or run.get("phase") or status)) detail = normalize_inline_whitespace(str(payload.get("detail") or "")) run_id = normalize_inline_whitespace(str(payload.get("run_id") or run.get("run_id") or "")) if not status and not phase: return None label = normalize_inline_whitespace(command.replace("-", " ")) line = f"{label} {phase or status}".strip() if detail: line += f" — {detail}" if run_id: line += f" (run_id {run_id})" return line or None def progress_payload_signature(command: str, payload: dict[str, object]) -> tuple[str, str, str, str, str]: run = cli_run_payload(payload) status = normalize_inline_whitespace(str(payload.get("status") or run.get("status") or "")) phase = normalize_inline_whitespace(str(payload.get("phase") or run.get("phase") or status)) run_id = normalize_inline_whitespace(str(payload.get("run_id") or run.get("run_id") or "")) next_commands = payload.get("next_recommended_commands") if not isinstance(next_commands, list): next_commands = run.get("next_recommended_commands") first_next_command = "" if isinstance(next_commands, list) and next_commands: first_next_command = normalize_inline_whitespace(str(next_commands[0] or "")) return (command, run_id, status, phase, first_next_command) class LongTaskProgressPrinter: def __init__(self, command: str, *, interval_seconds: float = LONG_TASK_PROGRESS_INTERVAL_SECONDS) -> None: self.command = command self.interval_seconds = max(1.0, float(interval_seconds)) self.last_emit_at = 0.0 self.last_signature: tuple[str, str, str, str, str] | None = None def emit(self, payload: dict[str, object], force: bool = False) -> None: line = render_progress_line(self.command, payload) if not line: return signature = progress_payload_signature(self.command, payload) now = time.perf_counter() should_emit = ( force or self.last_signature is None or signature != self.last_signature or (now - self.last_emit_at) >= self.interval_seconds ) if not should_emit: return print(f"[progress] {line}", file=sys.stderr, flush=True) self.last_emit_at = now self.last_signature = signature def emit_cli_payload(command: str, payload: dict[str, object], *, verbose: bool = False) -> int: benchmark_mark("prepare_payload_begin", command=command, verbose=verbose) prepared_payload = prepare_cli_payload(command, payload, verbose=verbose) benchmark_mark("prepare_payload_done") if CLI_OUTPUT_MODE == "human": rendered_human = render_cli_human_output(command, prepared_payload) if rendered_human is not None: sys.stdout.write(rendered_human + "\n") sys.stdout.flush() benchmark_mark("stdout_written") benchmark_emit(command=command, verbose=verbose) return 0 serialized = json.dumps(prepared_payload, indent=2, sort_keys=True) benchmark_mark("json_serialized", bytes=len(serialized.encode("utf-8"))) sys.stdout.write(serialized + "\n") sys.stdout.flush() benchmark_mark("stdout_written") benchmark_emit(command=command, verbose=verbose) return 0 def document_chunk_rows(connection: sqlite3.Connection, document_id: int) -> list[sqlite3.Row]: return connection.execute( """ SELECT id, document_id, chunk_index, char_start, char_end, token_estimate, text_content FROM document_chunks WHERE document_id = ? ORDER BY chunk_index ASC """, (document_id,), ).fetchall() def reconstruct_document_text_prefix(chunk_rows: list[sqlite3.Row], max_chars: int) -> str | None: if not chunk_rows or max_chars <= 0: return None parts: list[str] = [] current_end = 0 total_length = 0 for row in chunk_rows: chunk_text = str(row["text_content"] or "") chunk_start = int(row["char_start"]) if not chunk_text: continue overlap = max(0, current_end - chunk_start) append_text = chunk_text[overlap:] if not append_text: continue remaining = max_chars - total_length if remaining <= 0: break parts.append(append_text[:remaining]) total_length += len(parts[-1]) current_end = max(current_end, chunk_start + len(chunk_text)) if total_length >= max_chars: break combined = "".join(parts) return combined or None def document_overview_payload( paths: dict[str, Path], connection: sqlite3.Connection, row: sqlite3.Row, *, occurrence_row: sqlite3.Row | None = None, include_parent_context: bool = True, include_attachment_context: bool = True, ) -> dict[str, object]: source_row = occurrence_row or row custodian_values = document_custodian_values_from_row(row) custodian_text = ", ".join(custodian_values) if custodian_values else None payload: dict[str, object] = { "document_id": int(row["id"]), "control_number": row["control_number"], "conversation_id": row["conversation_id"], "dataset_id": row["dataset_id"], "parent_document_id": row["parent_document_id"], "child_document_kind": row["child_document_kind"], "root_message_key": row["root_message_key"], "source_kind": source_row["source_kind"], "source_rel_path": source_row["source_rel_path"], "source_item_id": source_row["source_item_id"], "source_folder_path": source_row["source_folder_path"], "production_id": source_row["production_id"], **document_path_payload(paths, connection, row, occurrence_row=occurrence_row), "file_name": source_row["file_name"], "file_type": source_row["file_type"], "custodian": custodian_text, "custodians": custodian_values, "metadata": { "author": row["author"], "begin_attachment": source_row["begin_attachment"], "begin_bates": source_row["begin_bates"], "child_document_kind": row["child_document_kind"], "content_type": row["content_type"], "conversation_id": row["conversation_id"], "custodian": custodian_text, "custodians": custodian_values, "dataset_id": row["dataset_id"], "date_created": row["date_created"], "date_modified": row["date_modified"], "end_attachment": source_row["end_attachment"], "end_bates": source_row["end_bates"], "page_count": row["page_count"], "participants": row["participants"], "recipients": row["recipients"], "root_message_key": row["root_message_key"], "source_kind": source_row["source_kind"], "source_rel_path": source_row["source_rel_path"], "source_item_id": source_row["source_item_id"], "source_folder_path": source_row["source_folder_path"], "subject": row["subject"], "title": row["title"], "updated_at": row["updated_at"], }, "manual_field_locks": normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN]), } dataset_memberships = fetch_document_dataset_memberships(connection, [row]).get(int(row["id"]), {"ids": [], "names": []}) dataset_ids = [int(dataset_id) for dataset_id in dataset_memberships["ids"]] dataset_names = [str(dataset_name) for dataset_name in dataset_memberships["names"]] payload["dataset_ids"] = dataset_ids payload["dataset_names"] = dataset_names if len(dataset_ids) == 1: payload["dataset_id"] = dataset_ids[0] payload["metadata"]["dataset_id"] = dataset_ids[0] else: payload["dataset_id"] = None payload["metadata"]["dataset_id"] = None if len(dataset_names) == 1: payload["dataset_name"] = dataset_names[0] payload["metadata"]["dataset_name"] = dataset_names[0] if row["production_id"] is not None: production_name = fetch_production_names(connection, [row]).get(int(row["production_id"])) payload["production_name"] = production_name payload["metadata"]["production_name"] = production_name if include_attachment_context and row["parent_document_id"] is None: attachments = fetch_attachment_summaries(connection, paths, [int(row["id"])]).get(int(row["id"]), []) payload["attachment_count"] = len(attachments) payload["attachments"] = attachments child_documents = fetch_child_document_summaries(connection, paths, [int(row["id"])]).get(int(row["id"]), []) payload["child_document_count"] = len(child_documents) payload["child_documents"] = child_documents if include_parent_context and row["parent_document_id"] is not None: payload["parent"] = fetch_parent_summaries(connection, [row]).get(int(row["parent_document_id"])) return payload def build_chunk_citation_payload( row: sqlite3.Row, occurrence_row: sqlite3.Row | None = None, *, preview_rel_path: str, preview_abs_path: str, chunk_index: int, char_start: int, char_end: int, snippet: str, ) -> dict[str, object]: source_row = occurrence_row or row return { "document_id": int(row["id"]), "control_number": row["control_number"], "file_name": source_row["file_name"], "chunk_index": chunk_index, "char_start": char_start, "char_end": char_end, "snippet": snippet, "preview_rel_path": preview_rel_path, "preview_abs_path": preview_abs_path, } def filter_operators_for_field_type(field_type: str) -> list[str]: return sql_filter_operator_names_for_field_type(field_type) def catalog_description_for_field(field_name: str, *, source: str, instruction: object = None) -> str: if source == "builtin": return BUILTIN_FIELD_DESCRIPTIONS.get(field_name, normalize_whitespace(field_name.replace("_", " "))) if source == "virtual": return VIRTUAL_FIELD_DESCRIPTIONS.get(field_name, normalize_whitespace(field_name.replace("_", " "))) if isinstance(instruction, str) and normalize_whitespace(instruction): return normalize_whitespace(instruction) return normalize_whitespace(field_name.replace("_", " ")) def catalog_field_entry( field_name: str, field_type: str, *, source: str, instruction: object = None, displayable: bool | None = None, ) -> dict[str, object]: return { "name": field_name, "type": field_type, "description": catalog_description_for_field(field_name, source=source, instruction=instruction), "filter_operators": filter_operators_for_field_type(field_type), "sortable": source != "virtual", "aggregatable": source != "virtual" or field_name in AGGREGATABLE_VIRTUAL_FIELDS, "displayable": displayable if displayable is not None else (source != "virtual" or field_name in DISPLAYABLE_VIRTUAL_FIELDS), "date_granularities": ["year", "quarter", "month", "week"] if field_type == "date" else [], } def search_bates( connection: sqlite3.Connection, query_begin: str, query_end: str, clauses: list[str], params: list[object], ) -> dict[int, dict[str, object]]: rows = connection.execute( f""" SELECT * FROM documents d WHERE {' AND '.join(clauses)} """, params, ).fetchall() single_value = query_begin == query_end matches: dict[int, dict[str, object]] = {} for row in rows: row_begin = row["begin_bates"] or row["control_number"] row_end = row["end_bates"] or row["control_number"] rank: float | None = None if single_value: if row["control_number"] == query_begin or row["begin_bates"] == query_begin or row["end_bates"] == query_begin: rank = 0.0 elif bates_inclusive_contains(row_begin, row_end, query_begin): rank = 1.0 else: if bates_ranges_overlap(row_begin, row_end, query_begin, query_end): rank = 2.0 if rank is None: continue document_id = int(row["id"]) matches[document_id] = { "row": row, "rank": rank, "snippet": metadata_snippet(row), "bates_sort_key": bates_sort_key(row_begin), } return matches def search_fts( connection: sqlite3.Connection, query: str, clauses: list[str], params: list[object], ) -> dict[int, dict[str, object]]: query_value = query.strip() if not query_value: return {} where_clause = " AND ".join(clauses) chunk_sql = f""" SELECT d.*, dc.text_content AS snippet_source, bm25(chunks_fts) AS rank FROM chunks_fts JOIN document_chunks dc ON dc.id = CAST(chunks_fts.chunk_id AS INTEGER) JOIN documents d ON d.id = dc.document_id WHERE chunks_fts MATCH ? AND {where_clause} """ metadata_sql = f""" SELECT d.*, NULL AS snippet_source, bm25(documents_fts) AS rank FROM documents_fts JOIN documents d ON d.id = CAST(documents_fts.document_id AS INTEGER) WHERE documents_fts MATCH ? AND {where_clause} """ matches: dict[int, dict[str, object]] = {} for sql, source in ((chunk_sql, "chunk"), (metadata_sql, "metadata")): try: rows = connection.execute(sql, [query_value, *params]).fetchall() except sqlite3.OperationalError: rows = connection.execute(sql, [f'"{query_value}"', *params]).fetchall() for row in rows: document_id = int(row["id"]) rank = float(row["rank"]) existing = matches.get(document_id) if existing is None or rank < float(existing["rank"]): matches[document_id] = { "row": row, "rank": rank, "snippet": make_snippet( row["snippet_source"] if source == "chunk" and row["snippet_source"] else metadata_snippet(row), query_value, ), } return matches def search_browse( connection: sqlite3.Connection, clauses: list[str], params: list[object], ) -> dict[int, dict[str, object]]: rows = connection.execute( f""" SELECT * FROM documents d WHERE {' AND '.join(clauses)} """, params, ).fetchall() return { int(row["id"]): { "row": row, "rank": None, "snippet": metadata_snippet(row), } for row in rows } def coerce_sort_value(value: object) -> tuple[int, object]: if value is None: return (1, "") if isinstance(value, str): return (0, value.lower()) return (0, value) LISTING_QUERY_PREFIXES = ( "show ", "show me ", "list ", "display ", "browse ", "open ", "get ", "get me ", "give me ", ) LISTING_QUERY_NOUNS = { "attachment", "attachments", "conversation", "conversations", "document", "documents", "email", "emails", "file", "files", "message", "messages", "result", "results", "thread", "threads", } LISTING_QUERY_SCOPE_TERMS = { "about", "for", "from", "in", "where", "with", "within", } THREAD_QUERY_TERMS = {"thread", "threads", "conversation", "conversations"} def keyword_query_uses_relevance_scoring(query: str) -> bool: normalized_query = normalize_inline_whitespace(query) if not normalized_query: return False lowered_query = normalized_query.lower() tokens = lowered_query.split() if not tokens: return False if any(lowered_query.startswith(prefix) for prefix in LISTING_QUERY_PREFIXES): return False if " -- " in normalized_query: return False if tokens[-1] in THREAD_QUERY_TERMS and len(tokens) > 1: return False if any(noun in tokens for noun in LISTING_QUERY_NOUNS) and any( scope_term in tokens for scope_term in LISTING_QUERY_SCOPE_TERMS ): return False if any(term in tokens for term in THREAD_QUERY_TERMS) and len(tokens) > 1: return False return True def sort_search_results( results: list[dict[str, object]], sort_field: str | None, order: str | None, query: str, ) -> list[dict[str, object]]: query_present = bool(query.strip()) stable_results = sorted(results, key=lambda item: item["id"]) if query_present and (sort_field is None or sort_field == "relevance"): if keyword_query_uses_relevance_scoring(query): stable_results = stable_sort_results_by_field(stable_results, "date_created", "desc") return sorted(stable_results, key=lambda item: (item["rank"] is None, item["rank"])) return stable_sort_results_by_field(stable_results, "date_created", "desc") field_name = sort_field or "date_created" normalized_order = (order or "desc").lower() return stable_sort_results_by_field(stable_results, field_name, normalized_order) def resolve_document_search( connection: sqlite3.Connection, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, ) -> dict[str, object]: filter_summary, clauses, params = build_search_filters(connection, raw_filters) normalized_sort_field = sort_field uses_relevance_scoring = keyword_query_uses_relevance_scoring(query) bates_query_begin, bates_query_end = parse_bates_query(query) is_bates_query = bates_query_begin is not None and bates_query_end is not None if sort_field == "relevance" and not query.strip(): raise RetrieverError("Sort 'relevance' requires a non-empty query.") if sort_field and sort_field != "relevance": sort_field_def = resolve_field_definition(connection, sort_field) if sort_field_def.get("source") == "virtual": raise RetrieverError(f"Cannot sort by virtual filter field: {sort_field}") normalized_sort_field = sort_field_def["field_name"] if is_bates_query: matches = search_bates(connection, str(bates_query_begin), str(bates_query_end), clauses, params) elif query.strip(): matches = search_fts(connection, query, clauses, params) else: matches = search_browse(connection, clauses, params) results = [ { "id": document_id, "rank": match["rank"], "snippet": match["snippet"], "bates_sort_key": match.get("bates_sort_key"), "row": match["row"], } for document_id, match in matches.items() ] if is_bates_query and normalized_sort_field is None and order is None: sorted_results = sorted( sorted(results, key=lambda item: item["id"]), key=lambda item: ( item["rank"] is None, item["rank"], item.get("bates_sort_key") or (1, "", 0, ""), ), ) else: sorted_results = sort_search_results(results, normalized_sort_field, order, query) default_sort = "date_created" default_order = "desc" if is_bates_query and query.strip(): default_sort = "bates" default_order = "asc" elif query.strip() and uses_relevance_scoring: default_sort = "relevance" default_order = "asc" return { "query": query, "filters": filter_summary, "sort": normalized_sort_field or default_sort, "order": (order or default_order).lower(), "sort_spec": f"{normalized_sort_field or default_sort} {(order or default_order).lower()}", "results": sorted_results, } def sql_sort_terms_for_field( connection: sqlite3.Connection, field_name: str, order: str, *, alias: str, ) -> list[str]: normalized_order = "DESC" if normalize_inline_whitespace(order).lower() == "desc" else "ASC" if field_name == "id": return [f"{alias}.id {normalized_order}"] field_def = resolve_field_definition(connection, field_name) if field_def.get("source") == "virtual": raise RetrieverError(f"Cannot sort by virtual filter field: {field_name}") canonical_name = str(field_def["field_name"]) column_expr = f"{alias}.{quote_identifier(canonical_name)}" field_type = str(field_def["field_type"]) if field_type == "date": normalized_expr = f"datetime({column_expr})" return [ f"CASE WHEN {column_expr} IS NULL OR {normalized_expr} IS NULL THEN 1 ELSE 0 END ASC", f"{normalized_expr} {normalized_order}", ] if field_type == "text": return [ f"CASE WHEN {column_expr} IS NULL THEN 1 ELSE 0 END ASC", f"LOWER({column_expr}) {normalized_order}", ] return [ f"CASE WHEN {column_expr} IS NULL THEN 1 ELSE 0 END ASC", f"{column_expr} {normalized_order}", ] def sql_order_by_for_sort_specs( connection: sqlite3.Connection, sort_specs: list[tuple[str, str]], *, alias: str, ) -> str: effective_specs = list(sort_specs) if not any(field_name == "id" for field_name, _ in effective_specs): effective_specs.append(("id", "asc")) terms: list[str] = [] for field_name, direction in effective_specs: terms.extend(sql_sort_terms_for_field(connection, field_name, direction, alias=alias)) return ", ".join(terms) def sql_relevance_order_by( connection: sqlite3.Connection, *, row_alias: str, rank_expr: str, ) -> str: return ", ".join( [ f"CASE WHEN {rank_expr} IS NULL THEN 1 ELSE 0 END ASC", f"{rank_expr} ASC", sql_order_by_for_sort_specs(connection, [("date_created", "desc")], alias=row_alias), ] ) def sql_bates_order_by(*, row_alias: str, prioritize_rank: bool, id_alias: str | None = None) -> str: terms: list[str] = [] if prioritize_rank: terms.extend( [ f"CASE WHEN {row_alias}.rank IS NULL THEN 1 ELSE 0 END ASC", f"{row_alias}.rank ASC", ] ) terms.extend( [ f"CASE WHEN {row_alias}.bates_sort_value IS NULL THEN 1 ELSE 0 END ASC", f"{row_alias}.bates_sort_value ASC", f"{id_alias or row_alias}.id ASC", ] ) return ", ".join(terms) def search_browse_page( connection: sqlite3.Connection, clauses: list[str], params: list[object], *, limit: int, offset: int, order_by_sql: str, ) -> dict[str, object]: where_clause = " AND ".join(clauses) count_row = connection.execute( f""" SELECT COUNT(*) AS total_hits FROM documents d WHERE {where_clause} """, params, ).fetchone() total_hits = int(count_row["total_hits"] or 0) if count_row is not None else 0 rows = connection.execute( f""" SELECT d.* FROM documents d WHERE {where_clause} ORDER BY {order_by_sql} LIMIT ? OFFSET ? """, [*params, limit, offset], ).fetchall() return { "total_hits": total_hits, "results": [ { "id": int(row["id"]), "rank": None, "snippet": metadata_snippet(row), "row": row, } for row in rows ], } def search_bates_page( connection: sqlite3.Connection, query_begin: str, query_end: str, clauses: list[str], params: list[object], *, limit: int, offset: int, order_by_sql: str, ) -> dict[str, object]: where_clause = " AND ".join(clauses) range_begin_expr = "COALESCE(d.begin_bates, d.control_number)" range_end_expr = "COALESCE(d.end_bates, d.control_number)" single_value = query_begin == query_end if single_value: rank_sql = "CASE WHEN d.control_number = ? OR d.begin_bates = ? OR d.end_bates = ? THEN 0.0 ELSE 1.0 END" rank_params: list[object] = [query_begin, query_begin, query_begin] match_clause = ( "d.control_number = ? OR d.begin_bates = ? OR d.end_bates = ? " f"OR ({range_begin_expr} <= ? AND {range_end_expr} >= ?)" ) match_params: list[object] = [query_begin, query_begin, query_begin, query_begin, query_begin] else: rank_sql = "2.0" rank_params = [] match_clause = f"{range_begin_expr} <= ? AND {range_end_expr} >= ?" match_params = [query_end, query_begin] cte_sql = f""" WITH bates_matches AS ( SELECT d.*, {rank_sql} AS rank, {range_begin_expr} AS bates_sort_value FROM documents d WHERE {where_clause} AND ({match_clause}) ) """ query_params = [*rank_params, *params, *match_params] count_row = connection.execute( f""" {cte_sql} SELECT COUNT(*) AS total_hits FROM bates_matches """, query_params, ).fetchone() total_hits = int(count_row["total_hits"] or 0) if count_row is not None else 0 rows = connection.execute( f""" {cte_sql} SELECT * FROM bates_matches bm ORDER BY {order_by_sql} LIMIT ? OFFSET ? """, [*query_params, limit, offset], ).fetchall() return { "total_hits": total_hits, "results": [ { "id": int(row["id"]), "rank": float(row["rank"]) if row["rank"] is not None else None, "snippet": metadata_snippet(row), "bates_sort_key": bates_sort_key(row["bates_sort_value"] or row["control_number"]), "row": row, } for row in rows ], } def search_fts_page( connection: sqlite3.Connection, query: str, clauses: list[str], params: list[object], *, limit: int, offset: int, order_by_sql: str, ) -> dict[str, object]: query_value = query.strip() if not query_value: return {"total_hits": 0, "results": []} where_clause = " AND ".join(clauses) cte_sql = f""" WITH chunk_matches AS ( SELECT d.id AS document_id, dc.text_content AS snippet_source, bm25(chunks_fts) AS rank, 0 AS source_priority FROM chunks_fts JOIN document_chunks dc ON dc.id = CAST(chunks_fts.chunk_id AS INTEGER) JOIN documents d ON d.id = dc.document_id WHERE chunks_fts MATCH ? AND {where_clause} ), metadata_matches AS ( SELECT d.id AS document_id, NULL AS snippet_source, bm25(documents_fts) AS rank, 1 AS source_priority FROM documents_fts JOIN documents d ON d.id = CAST(documents_fts.document_id AS INTEGER) WHERE documents_fts MATCH ? AND {where_clause} ), all_matches AS ( SELECT * FROM chunk_matches UNION ALL SELECT * FROM metadata_matches ), ranked_matches AS ( SELECT document_id, snippet_source, rank, source_priority, ROW_NUMBER() OVER ( PARTITION BY document_id ORDER BY rank ASC, source_priority ASC, document_id ASC ) AS row_number FROM all_matches ), best_matches AS ( SELECT document_id, snippet_source, rank FROM ranked_matches WHERE row_number = 1 ) """ def fts_params(fts_query: str) -> list[object]: return [fts_query, *params, fts_query, *params] count_sql = f""" {cte_sql} SELECT COUNT(*) AS total_hits FROM best_matches """ effective_query = query_value try: count_row = connection.execute(count_sql, fts_params(effective_query)).fetchone() except sqlite3.OperationalError: effective_query = f'"{query_value}"' count_row = connection.execute(count_sql, fts_params(effective_query)).fetchone() total_hits = int(count_row["total_hits"] or 0) if count_row is not None else 0 rows = connection.execute( f""" {cte_sql} SELECT d.*, bm.rank AS rank, bm.snippet_source AS snippet_source FROM best_matches bm JOIN documents d ON d.id = bm.document_id ORDER BY {order_by_sql} LIMIT ? OFFSET ? """, [*fts_params(effective_query), limit, offset], ).fetchall() results: list[dict[str, object]] = [] for row in rows: snippet_source = row["snippet_source"] if row["snippet_source"] else metadata_snippet(row) results.append( { "id": int(row["id"]), "rank": float(row["rank"]) if row["rank"] is not None else None, "snippet": make_snippet(snippet_source, query_value), "row": row, } ) return {"total_hits": total_hits, "results": results} def resolve_paged_document_search( connection: sqlite3.Connection, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, *, page: int, per_page: int, ) -> dict[str, object]: filter_summary, clauses, params = build_search_filters(connection, raw_filters) normalized_sort_field = sort_field uses_relevance_scoring = keyword_query_uses_relevance_scoring(query) bates_query_begin, bates_query_end = parse_bates_query(query) is_bates_query = bates_query_begin is not None and bates_query_end is not None if sort_field == "relevance" and not query.strip(): raise RetrieverError("Sort 'relevance' requires a non-empty query.") if sort_field and sort_field != "relevance": sort_field_def = resolve_field_definition(connection, sort_field) if sort_field_def.get("source") == "virtual": raise RetrieverError(f"Cannot sort by virtual filter field: {sort_field}") normalized_sort_field = str(sort_field_def["field_name"]) offset = max(0, (page - 1) * per_page) if is_bates_query: if normalized_sort_field is None and order is None: selection_page = search_bates_page( connection, str(bates_query_begin), str(bates_query_end), clauses, params, limit=per_page, offset=offset, order_by_sql=sql_bates_order_by(row_alias="bm", prioritize_rank=True), ) elif normalized_sort_field == "relevance": selection_page = search_bates_page( connection, str(bates_query_begin), str(bates_query_end), clauses, params, limit=per_page, offset=offset, order_by_sql=sql_relevance_order_by(connection, row_alias="bm", rank_expr="bm.rank"), ) else: effective_field = normalized_sort_field or "date_created" effective_order = (order or "desc").lower() selection_page = search_bates_page( connection, str(bates_query_begin), str(bates_query_end), clauses, params, limit=per_page, offset=offset, order_by_sql=sql_order_by_for_sort_specs(connection, [(effective_field, effective_order)], alias="bm"), ) elif query.strip(): if normalized_sort_field is None or normalized_sort_field == "relevance": if uses_relevance_scoring: order_by_sql = sql_relevance_order_by(connection, row_alias="d", rank_expr="bm.rank") else: order_by_sql = sql_order_by_for_sort_specs(connection, [("date_created", "desc")], alias="d") else: order_by_sql = sql_order_by_for_sort_specs( connection, [(normalized_sort_field, (order or "desc").lower())], alias="d", ) selection_page = search_fts_page( connection, query, clauses, params, limit=per_page, offset=offset, order_by_sql=order_by_sql, ) else: selection_page = search_browse_page( connection, clauses, params, limit=per_page, offset=offset, order_by_sql=sql_order_by_for_sort_specs( connection, [(normalized_sort_field or "date_created", (order or "desc").lower())], alias="d", ), ) default_sort = "date_created" default_order = "desc" if is_bates_query and query.strip(): default_sort = "bates" default_order = "asc" elif query.strip() and uses_relevance_scoring: default_sort = "relevance" default_order = "asc" return { "query": query, "filters": filter_summary, "sort": normalized_sort_field or default_sort, "order": (order or default_order).lower(), "sort_spec": f"{normalized_sort_field or default_sort} {(order or default_order).lower()}", "results": selection_page["results"], "total_hits": int(selection_page["total_hits"]), } def search( root: Path, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, page: int, per_page: int | None, raw_columns: str | None = None, mode: str = "compose", *, compact_mode: bool = False, ) -> dict[str, object]: if page < 1: raise RetrieverError("Page must be >= 1.") paths = workspace_paths(root) ensure_layout(paths) if per_page is None: per_page = session_page_size(read_session_state(paths), browse_mode=BROWSE_MODE_DOCUMENTS) if per_page < 1: raise RetrieverError("per-page must be >= 1.") per_page = min(per_page, MAX_PAGE_SIZE) normalized_mode = normalize_search_mode(mode) compact_mode = bool(compact_mode) and normalized_mode == "compose" connection = connect_db(paths["db_path"]) try: benchmark_mark("schema_begin") apply_schema(connection, root) benchmark_mark("schema_done") selection = resolve_paged_document_search( connection, query, raw_filters, sort_field, order, page=page, per_page=per_page, ) benchmark_mark("query_done", total_hits=int(selection["total_hits"])) derived_scope = derive_search_scope(query, raw_filters) raw_column_list = parse_display_columns_argument(raw_columns) if raw_columns is not None else None display_column_defs, display_warnings, _ = resolve_display_column_definitions( connection, raw_column_list, drop_missing=False, ) occurrence_scope_clauses, occurrence_scope_params = build_occurrence_scope_filters(connection, raw_filters) results: list[dict[str, object]] = [] for match in selection["results"]: row = match["row"] if compact_mode: results.append( { "id": int(match["id"]), "control_number": row["control_number"], "dataset_id": row["dataset_id"], "parent_document_id": row["parent_document_id"], "production_id": row["production_id"], **document_path_payload(paths, connection, row, include_preview_targets=False), "file_name": row["file_name"], "file_type": row["file_type"], "snippet": str(match["snippet"]), "rank": match["rank"], "metadata": { key: row[key] for key in COMPACT_METADATA_FIELDS if key in row.keys() and payload_has_meaningful_value(row[key]) }, "row": row, } ) else: results.append( { "id": int(match["id"]), "control_number": row["control_number"], "conversation_id": row["conversation_id"], "dataset_id": row["dataset_id"], "parent_document_id": row["parent_document_id"], "child_document_kind": row["child_document_kind"], "root_message_key": row["root_message_key"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "source_folder_path": row["source_folder_path"], "production_id": row["production_id"], **document_path_payload(paths, connection, row), "file_name": row["file_name"], "file_type": row["file_type"], "snippet": str(match["snippet"]), "rank": match["rank"], "bates_sort_key": match.get("bates_sort_key"), "metadata": { "author": row["author"], "begin_attachment": row["begin_attachment"], "begin_bates": row["begin_bates"], "child_document_kind": row["child_document_kind"], "content_type": row["content_type"], "conversation_id": row["conversation_id"], "custodian": document_custodian_display_text_from_row(row), "custodians": document_custodian_values_from_row(row), "dataset_id": row["dataset_id"], "date_created": row["date_created"], "date_modified": row["date_modified"], "end_attachment": row["end_attachment"], "end_bates": row["end_bates"], "page_count": row["page_count"], "participants": row["participants"], "recipients": row["recipients"], "root_message_key": row["root_message_key"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "source_folder_path": row["source_folder_path"], "subject": row["subject"], "title": row["title"], "updated_at": row["updated_at"], }, "manual_field_locks": normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN]), "row": row, } ) benchmark_mark("full_matchset_built", page_matches=len(results)) paged_results = results total_hits = int(selection["total_hits"]) total_pages = max(1, (total_hits + per_page - 1) // per_page) paged_rows = [item["row"] for item in paged_results] preferred_occurrences = preferred_occurrences_by_document( connection, [int(row["id"]) for row in paged_rows], occurrence_scope_clauses, occurrence_scope_params, ) paged_parent_ids = [int(row["id"]) for row in paged_rows if row["parent_document_id"] is None] production_names = fetch_production_names(connection, paged_rows) dataset_memberships = fetch_document_dataset_memberships(connection, paged_rows) attachment_counts: dict[int, int] = {} child_document_counts: dict[int, int] = {} attachment_summaries: dict[int, list[dict[str, object]]] = {} child_document_summaries: dict[int, list[dict[str, object]]] = {} if compact_mode: attachment_counts = fetch_attachment_counts(connection, paged_parent_ids) child_document_counts = fetch_child_document_counts(connection, paged_parent_ids) else: attachment_summaries = fetch_attachment_summaries( connection, paths, paged_parent_ids, ) child_document_summaries = fetch_child_document_summaries( connection, paths, paged_parent_ids, ) parent_summaries = fetch_parent_summaries( connection, [row for row in paged_rows if row["parent_document_id"] is not None], ) for item in paged_results: row = item["row"] occurrence_row = preferred_occurrences.get(int(row["id"])) source_row = occurrence_row or row custodian_values = document_custodian_values_from_row(row) custodian_text = ", ".join(custodian_values) if custodian_values else None path_payload = document_path_payload( paths, connection, row, occurrence_row=occurrence_row, include_preview_targets=not compact_mode, ) item["rel_path"] = path_payload["rel_path"] item["abs_path"] = path_payload["abs_path"] item["preview_rel_path"] = path_payload["preview_rel_path"] item["preview_abs_path"] = path_payload["preview_abs_path"] item["preview_file_rel_path"] = path_payload["preview_file_rel_path"] item["preview_file_abs_path"] = path_payload["preview_file_abs_path"] item["preview_target_fragment"] = path_payload["preview_target_fragment"] if not compact_mode: item["preview_targets"] = path_payload["preview_targets"] item["file_name"] = source_row["file_name"] item["file_type"] = source_row["file_type"] item["custodian"] = custodian_text item["custodians"] = custodian_values if compact_mode: item["metadata"]["custodian"] = custodian_text item["metadata"]["custodians"] = custodian_values else: item["source_kind"] = source_row["source_kind"] item["source_rel_path"] = source_row["source_rel_path"] item["source_item_id"] = source_row["source_item_id"] item["source_folder_path"] = source_row["source_folder_path"] item["production_id"] = source_row["production_id"] item["metadata"]["begin_attachment"] = source_row["begin_attachment"] item["metadata"]["begin_bates"] = source_row["begin_bates"] item["metadata"]["custodian"] = custodian_text item["metadata"]["custodians"] = custodian_values item["metadata"]["end_attachment"] = source_row["end_attachment"] item["metadata"]["end_bates"] = source_row["end_bates"] item["metadata"]["source_kind"] = source_row["source_kind"] item["metadata"]["source_rel_path"] = source_row["source_rel_path"] item["metadata"]["source_item_id"] = source_row["source_item_id"] item["metadata"]["source_folder_path"] = source_row["source_folder_path"] memberships = dataset_memberships.get(int(row["id"]), {"ids": [], "names": []}) dataset_ids = [int(dataset_id) for dataset_id in memberships["ids"]] dataset_names = [str(dataset_name) for dataset_name in memberships["names"]] item["dataset_ids"] = dataset_ids item["dataset_names"] = dataset_names item["metadata"]["dataset_ids"] = dataset_ids item["metadata"]["dataset_names"] = dataset_names if len(dataset_ids) == 1: item["dataset_id"] = dataset_ids[0] item["metadata"]["dataset_id"] = dataset_ids[0] else: item["dataset_id"] = None item["metadata"]["dataset_id"] = None if len(dataset_names) == 1: item["dataset_name"] = dataset_names[0] item["metadata"]["dataset_name"] = dataset_names[0] if source_row["production_id"] is not None: item["production_name"] = production_names.get(int(source_row["production_id"])) item["metadata"]["production_name"] = production_names.get(int(source_row["production_id"])) if row["parent_document_id"] is None: if compact_mode: attachment_count = attachment_counts.get(int(row["id"]), 0) if attachment_count > 0: item["attachment_count"] = attachment_count child_document_count = child_document_counts.get(int(row["id"]), 0) if child_document_count > 0: item["child_document_count"] = child_document_count else: attachments = attachment_summaries.get(int(row["id"]), []) item["attachment_count"] = len(attachments) item["attachments"] = attachments child_documents = child_document_summaries.get(int(row["id"]), []) item["child_document_count"] = len(child_documents) item["child_documents"] = child_documents else: item["parent"] = parent_summaries.get(int(row["parent_document_id"])) item["display_values"] = build_search_result_display_values(row, item, display_column_defs) item.pop("bates_sort_key", None) item.pop("row", None) benchmark_mark("page_enriched", page_size=len(paged_results), total_hits=total_hits) payload = { "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], "sort_spec": selection["sort_spec"], "browse_mode": BROWSE_MODE_DOCUMENTS, "page": page, "per_page": per_page, "total_hits": total_hits, "total_pages": total_pages, "results": paged_results, "scope": derived_scope, "display": build_display_payload(display_column_defs, per_page), "header": build_search_header_payload(derived_scope, { "sort": selection["sort"], "order": selection["order"], "sort_spec": selection["sort_spec"], "page": page, "per_page": per_page, "total_hits": total_hits, "total_pages": total_pages, }), } if display_warnings: payload["warnings"] = display_warnings payload["rendered_markdown"] = render_search_markdown(payload, display_column_defs) if normalized_mode == "view": explicit_sort_specs = None if sort_field: explicit_sort_specs = [(str(selection["sort"]), str(selection["order"]))] persist_direct_view_search_result( root, payload, display_column_defs, sort_specs=explicit_sort_specs, browse_mode=BROWSE_MODE_DOCUMENTS, ) return payload finally: connection.close() def format_scope_bates_value(bates_scope: object) -> str: if not isinstance(bates_scope, dict): return "" begin = normalize_inline_whitespace(str(bates_scope.get("begin") or "")) end = normalize_inline_whitespace(str(bates_scope.get("end") or "")) if not begin or not end: return "" return begin if begin == end else f"{begin}-{end}" def truncate_scope_header_value(value: object, *, max_chars: int = 200) -> str: normalized = normalize_inline_whitespace(str(value or "")) if len(normalized) <= max_chars: return normalized return normalized[: max_chars - 1] + "…" def format_scope_header(scope: dict[str, object]) -> str: parts: list[str] = [] if isinstance(scope.get("keyword"), str) and scope["keyword"].strip(): parts.append(f"keyword={scope['keyword']!r}") bates_value = format_scope_bates_value(scope.get("bates")) if bates_value: parts.append(f"bates={bates_value}") if isinstance(scope.get("filter"), str) and scope["filter"].strip(): parts.append(f"filter={truncate_scope_header_value(scope['filter'])}") dataset_entries = coerce_scope_dataset_entries(scope.get("dataset")) if dataset_entries: dataset_names = ", ".join(entry["name"] for entry in dataset_entries) parts.append(f"dataset={truncate_scope_header_value(dataset_names)}") if scope.get("from_run_id") is not None: parts.append(f"from_run_id={scope['from_run_id']}") return "Scope: (none)" if not parts else "Scope: " + ", ".join(parts) def derive_search_scope(query: str, raw_filters: object | None) -> dict[str, object]: scope: dict[str, object] = {} bates_begin, bates_end = parse_bates_query(query) if bates_begin is not None and bates_end is not None: scope["bates"] = {"begin": bates_begin, "end": bates_end} elif query.strip(): scope["keyword"] = query if raw_filters: expressions = normalize_sql_filter_expressions(raw_filters) if expressions: scope["filter"] = " AND ".join(f"({expression})" for expression in expressions) return scope def build_search_header_payload(scope: dict[str, object], payload: dict[str, object]) -> dict[str, str]: total_hits = int(payload.get("total_hits") or 0) page = int(payload.get("page") or 1) per_page = int(payload.get("per_page") or DEFAULT_PAGE_SIZE) start_index = 0 if total_hits == 0 else ((page - 1) * per_page) + 1 end_index = 0 if total_hits == 0 else min(total_hits, page * per_page) sort_summary = str(payload.get("sort_spec") or f"{payload.get('sort')} {payload.get('order')}") browse_mode = normalize_browse_mode(payload.get("browse_mode")) result_label = browse_mode_result_label(browse_mode) header: dict[str, str] = {} keyword = scope.get("keyword") if isinstance(scope.get("keyword"), str) else None if keyword and keyword.strip(): header["keyword"] = f"Keyword: {keyword!r}" bates_value = format_scope_bates_value(scope.get("bates")) if bates_value: header["bates"] = f"Bates: {bates_value}" filter_value = scope.get("filter") if isinstance(filter_value, str) and filter_value.strip(): header["filters"] = f"Active filters: {truncate_scope_header_value(filter_value)}" dataset_entries = coerce_scope_dataset_entries(scope.get("dataset")) if dataset_entries: dataset_names = ", ".join(entry["name"] for entry in dataset_entries) header["datasets"] = f"Datasets: {truncate_scope_header_value(dataset_names)}" if scope.get("from_run_id") is not None: header["from_run_id"] = f"From run: {scope['from_run_id']}" if not any(key in header for key in ("keyword", "bates", "filters", "datasets", "from_run_id")): header["scope"] = "Scope: (none)" header["sort"] = f"Sort: {sort_summary}" header["page"] = ( f"Page: {page} of {payload.get('total_pages')}" f" ({result_label} {start_index}-{end_index} of {total_hits})" ) return header CONVERSATION_FIELD_DEFINITIONS = { "conversation_type": { "field_name": "conversation_type", "field_type": "text", "source": "conversation", "displayable": "true", "sortable": "true", }, "title": { "field_name": "title", "field_type": "text", "source": "conversation", "displayable": "true", "sortable": "true", }, "participants": { "field_name": "participants", "field_type": "text", "source": "conversation", "displayable": "true", "sortable": "false", }, "first_activity": { "field_name": "first_activity", "field_type": "date", "source": "conversation", "displayable": "true", "sortable": "true", }, "last_activity": { "field_name": "last_activity", "field_type": "date", "source": "conversation", "displayable": "true", "sortable": "true", }, "document_count": { "field_name": "document_count", "field_type": "integer", "source": "conversation", "displayable": "true", "sortable": "true", }, "matching_document_count": { "field_name": "matching_document_count", "field_type": "integer", "source": "conversation", "displayable": "true", "sortable": "true", }, "source_kind": { "field_name": "source_kind", "field_type": "text", "source": "conversation", "displayable": "true", "sortable": "true", }, "dataset_name": { "field_name": "dataset_name", "field_type": "text", "source": "conversation", "displayable": "true", "sortable": "false", }, } ENTITY_FIELD_DEFINITIONS = { "id": { "field_name": "id", "field_type": "integer", "source": "entity", "displayable": "true", "sortable": "true", }, "label": { "field_name": "label", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "entity_type": { "field_name": "entity_type", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "entity_origin": { "field_name": "entity_origin", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "entity_status": { "field_name": "entity_status", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "canonical_status": { "field_name": "entity_status", "field_type": "text", "source": "entity", "displayable": "false", "sortable": "true", }, "display_name": { "field_name": "display_name", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "primary_email": { "field_name": "primary_email", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "primary_phone": { "field_name": "primary_phone", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "sort_name": { "field_name": "sort_name", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "true", }, "document_count": { "field_name": "document_count", "field_type": "integer", "source": "entity", "displayable": "true", "sortable": "true", }, "roles": { "field_name": "roles", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "false", }, "emails": { "field_name": "emails", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "false", }, "names": { "field_name": "names", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "false", }, "phones": { "field_name": "phones", "field_type": "text", "source": "entity", "displayable": "true", "sortable": "false", }, } def entity_field_definition(field_name: str) -> dict[str, str]: canonical_name = FIELD_NAME_ALIASES.get(field_name, field_name) field_def = ENTITY_FIELD_DEFINITIONS.get(canonical_name) if field_def is None: raise RetrieverError(f"Unknown field: {field_name}") return dict(field_def) def default_display_columns(browse_mode: str = BROWSE_MODE_DOCUMENTS) -> list[str]: effective_browse_mode = normalize_browse_mode(browse_mode) if effective_browse_mode == BROWSE_MODE_CONVERSATIONS: return list(DEFAULT_CONVERSATION_DISPLAY_COLUMNS) if effective_browse_mode == BROWSE_MODE_ENTITIES: return list(DEFAULT_ENTITY_DISPLAY_COLUMNS) return list(DEFAULT_DOCUMENT_DISPLAY_COLUMNS) def normalize_search_mode(mode: object | None) -> str: normalized = normalize_inline_whitespace(str(mode or "compose")).lower() if not normalized: return "compose" if normalized not in {"compose", "view"}: raise RetrieverError(f"Unknown search mode: {mode!r}. Expected 'compose' or 'view'.") return normalized def session_browse_mode(session_state: dict[str, object]) -> str: return normalize_browse_mode(session_state.get("browse_mode")) def session_display_state( session_state: dict[str, object], *, browse_mode: str | None = None, ) -> dict[str, object]: effective_browse_mode = normalize_browse_mode(browse_mode or session_browse_mode(session_state)) display = session_state.get("display") if not isinstance(display, dict): return {} branch = display.get(effective_browse_mode) return branch if isinstance(branch, dict) else {} def conversation_field_definition(field_name: str) -> dict[str, str]: canonical_name = FIELD_NAME_ALIASES.get(field_name, field_name) field_def = CONVERSATION_FIELD_DEFINITIONS.get(canonical_name) if field_def is None: raise RetrieverError(f"Unknown field: {field_name}") return dict(field_def) def resolve_browse_field_definition( connection: sqlite3.Connection, field_name: str, *, browse_mode: str, ) -> dict[str, str]: effective_browse_mode = normalize_browse_mode(browse_mode) if effective_browse_mode == BROWSE_MODE_CONVERSATIONS: return conversation_field_definition(field_name) if effective_browse_mode == BROWSE_MODE_ENTITIES: return entity_field_definition(field_name) return resolve_field_definition(connection, field_name) def known_browse_field_names(connection: sqlite3.Connection, *, browse_mode: str) -> list[str]: effective_browse_mode = normalize_browse_mode(browse_mode) if effective_browse_mode == BROWSE_MODE_CONVERSATIONS: return sorted(CONVERSATION_FIELD_DEFINITIONS) if effective_browse_mode == BROWSE_MODE_ENTITIES: return sorted(ENTITY_FIELD_DEFINITIONS) return known_logical_field_names(connection) def displayable_field_names( connection: sqlite3.Connection, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> list[str]: names: set[str] = set() for field_name in known_browse_field_names(connection, browse_mode=browse_mode): field_def = resolve_browse_field_definition(connection, field_name, browse_mode=browse_mode) if str(field_def.get("displayable") or "").lower() == "true": names.add(str(field_def["field_name"])) return sorted(names) def display_field_suggestions( connection: sqlite3.Connection, field_name: str, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> list[str]: return difflib.get_close_matches( field_name, displayable_field_names(connection, browse_mode=browse_mode), n=3, cutoff=0.45, ) def displayable_field_examples( connection: sqlite3.Connection, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, limit: int = 12, ) -> str: return ", ".join(displayable_field_names(connection, browse_mode=browse_mode)[:limit]) def default_display_column_definitions( connection: sqlite3.Connection, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> list[dict[str, str]]: definitions: list[dict[str, str]] = [] for column_name in default_display_columns(browse_mode): field_def = resolve_browse_field_definition(connection, column_name, browse_mode=browse_mode) definitions.append( { "name": str(field_def["field_name"]), "type": str(field_def["field_type"]), "source": str(field_def["source"]), } ) return definitions def resolve_display_column_definitions( connection: sqlite3.Connection, raw_columns: object, *, drop_missing: bool, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> tuple[list[dict[str, str]], list[str], bool]: raw_values = raw_columns if isinstance(raw_columns, list) and raw_columns else default_display_columns(browse_mode) warnings: list[str] = [] changed = False resolved_columns: list[dict[str, str]] = [] seen_names: set[str] = set() for raw_value in raw_values: normalized_name = normalize_inline_whitespace(str(raw_value or "")) if not normalized_name: changed = True continue try: field_def = resolve_browse_field_definition( connection, normalized_name, browse_mode=browse_mode, ) except RetrieverError as exc: if drop_missing: warnings.append( f"Column '{normalized_name}' no longer exists and has been removed from your display preferences." ) changed = True continue suggestions = display_field_suggestions(connection, normalized_name, browse_mode=browse_mode) suggestion_text = f" Did you mean: {', '.join(suggestions)}?" if suggestions else "" raise RetrieverError(f"Unknown column: {normalized_name}.{suggestion_text}") from exc canonical_name = str(field_def["field_name"]) if str(field_def.get("displayable") or "").lower() != "true": if drop_missing: warnings.append( f"Column '{canonical_name}' is no longer displayable and has been removed from your display preferences." ) changed = True continue raise RetrieverError( f"Field '{canonical_name}' is filter-only and cannot be displayed. " f"Displayable fields include: {displayable_field_examples(connection, browse_mode=browse_mode)}." ) if canonical_name in seen_names: changed = True continue if canonical_name != normalized_name: changed = True seen_names.add(canonical_name) resolved_columns.append( { "name": canonical_name, "type": str(field_def["field_type"]), "source": str(field_def["source"]), } ) if not resolved_columns: return default_display_column_definitions(connection, browse_mode=browse_mode), warnings, True return resolved_columns, warnings, changed def display_column_names(column_defs: list[dict[str, str]]) -> list[str]: return [column_def["name"] for column_def in column_defs] def build_display_payload(column_defs: list[dict[str, str]], page_size: int) -> dict[str, object]: return { "columns": display_column_names(column_defs), "page_size": page_size, } def persist_display_columns( paths: dict[str, Path], session_state: dict[str, object], column_defs: list[dict[str, str]], *, browse_mode: str, ) -> dict[str, object]: effective_browse_mode = normalize_browse_mode(browse_mode) display_root = session_state.get("display") if not isinstance(display_root, dict): display_root = {} display_state = session_display_state(session_state, browse_mode=effective_browse_mode) column_names = display_column_names(column_defs) if column_names == default_display_columns(effective_browse_mode): display_state.pop("columns", None) else: display_state["columns"] = column_names display_root[effective_browse_mode] = coerce_display_payload(display_state) session_state["display"] = coerce_mode_payloads(display_root, coerce_display_payload) return persist_session_state(paths, session_state) def persist_display_preferences( paths: dict[str, Path], session_state: dict[str, object], column_defs: list[dict[str, str]], page_size: int, *, browse_mode: str, ) -> dict[str, object]: effective_browse_mode = normalize_browse_mode(browse_mode) display_root = session_state.get("display") if not isinstance(display_root, dict): display_root = {} display_state = session_display_state(session_state, browse_mode=effective_browse_mode) column_names = display_column_names(column_defs) if column_names == default_display_columns(effective_browse_mode): display_state.pop("columns", None) else: display_state["columns"] = column_names if page_size == DEFAULT_PAGE_SIZE: display_state.pop("page_size", None) else: display_state["page_size"] = page_size display_root[effective_browse_mode] = coerce_display_payload(display_state) session_state["display"] = coerce_mode_payloads(display_root, coerce_display_payload) return persist_session_state(paths, session_state) def resolve_session_display_columns( connection: sqlite3.Connection, paths: dict[str, Path], session_state: dict[str, object], *, browse_mode: str, ) -> tuple[list[dict[str, str]], list[str], dict[str, object]]: configured_columns = session_display_state(session_state, browse_mode=browse_mode).get("columns") column_defs, warnings, changed = resolve_display_column_definitions( connection, configured_columns, drop_missing=True, browse_mode=browse_mode, ) if changed: session_state = persist_display_columns( paths, session_state, column_defs, browse_mode=browse_mode, ) return column_defs, warnings, session_state def parse_display_columns_argument(raw_text: str) -> list[str]: column_names = split_quoted_comma_values(raw_text) if not column_names: raise RetrieverError("Column list cannot be empty.") return column_names def parse_page_size_value(raw_value: str) -> int: try: page_size = int(normalize_inline_whitespace(raw_value)) except ValueError as exc: raise RetrieverError("Usage: /page-size ") from exc if page_size < 1 or page_size > MAX_PAGE_SIZE: raise RetrieverError(f"Page size must be between 1 and {MAX_PAGE_SIZE}.") return page_size def best_result_title(row: sqlite3.Row) -> str | None: for candidate in (row["title"], row["subject"], row["file_name"]): normalized = normalize_inline_whitespace(str(candidate or "")) if normalized: return normalized return None def search_result_display_value( row: sqlite3.Row, item: dict[str, object], field_name: str, field_type: str, ) -> object: if field_name == "title": return best_result_title(row) if field_name == "custodian": custodians = item.get("custodians") if isinstance(custodians, list): normalized_values = [normalize_inline_whitespace(str(value)) for value in custodians if normalize_inline_whitespace(str(value))] return ", ".join(normalized_values) or None value = item.get("custodian") return normalize_inline_whitespace(str(value)) or None if field_name == "dataset_name": dataset_names = item.get("dataset_names") if isinstance(dataset_names, list): normalized_names = [normalize_inline_whitespace(str(value)) for value in dataset_names if normalize_inline_whitespace(str(value))] return ", ".join(normalized_names) or None value = item.get("dataset_name") return normalize_inline_whitespace(str(value)) or None if field_name == "production_name": value = item.get("production_name") return normalize_inline_whitespace(str(value)) or None if field_name == "is_attachment": return "Yes" if row["parent_document_id"] is not None else "No" if field_name in item and field_name not in {"metadata", "display_values"}: value = item.get(field_name) elif field_name in row.keys(): value = row[field_name] else: metadata = item.get("metadata") value = metadata.get(field_name) if isinstance(metadata, dict) else None if field_type == "boolean": if value in (None, ""): return None return "Yes" if bool(value) else "No" if isinstance(value, list): normalized_values = [normalize_inline_whitespace(str(entry)) for entry in value if normalize_inline_whitespace(str(entry))] return ", ".join(normalized_values) or None if isinstance(value, str): return value or None return value def build_search_result_display_values( row: sqlite3.Row, item: dict[str, object], column_defs: list[dict[str, str]], ) -> dict[str, object]: return { column_def["name"]: search_result_display_value(row, item, column_def["name"], column_def["type"]) for column_def in column_defs } def build_summary_display_values( item: dict[str, object], column_defs: list[dict[str, str]], ) -> dict[str, object]: return { column_def["name"]: summary_display_value(item, column_def["name"], column_def["type"]) for column_def in column_defs } def best_summary_title(item: dict[str, object]) -> str | None: for candidate in ( item.get("title"), item.get("subject"), item.get("file_name"), item.get("control_number"), ): normalized = normalize_inline_whitespace(str(candidate or "")) if normalized: return normalized return None def summary_display_value(item: dict[str, object], field_name: str, field_type: str) -> object: if field_name == "title": return best_summary_title(item) if field_name == "custodian": custodians = item.get("custodians") if isinstance(custodians, list): normalized_values = [normalize_inline_whitespace(str(value)) for value in custodians if normalize_inline_whitespace(str(value))] return ", ".join(normalized_values) or None return normalize_inline_whitespace(str(item.get("custodian") or "")) or None if field_name == "dataset_name": dataset_names = item.get("dataset_names") if isinstance(dataset_names, list): normalized_names = [normalize_inline_whitespace(str(value)) for value in dataset_names if normalize_inline_whitespace(str(value))] return ", ".join(normalized_names) or None return normalize_inline_whitespace(str(item.get("dataset_name") or "")) or None if field_name == "production_name": return normalize_inline_whitespace(str(item.get("production_name") or "")) or None if field_name == "is_attachment": if item.get("parent_document_id") is not None: return "Yes" child_kind = normalize_inline_whitespace(str(item.get("child_document_kind") or "")) return "Yes" if child_kind == CHILD_DOCUMENT_KIND_ATTACHMENT else "No" if field_name in item: value = item.get(field_name) else: metadata = item.get("metadata") value = metadata.get(field_name) if isinstance(metadata, dict) else None if field_type == "boolean": if value in (None, ""): return None return "Yes" if bool(value) else "No" if isinstance(value, list): normalized_values = [normalize_inline_whitespace(str(entry)) for entry in value if normalize_inline_whitespace(str(entry))] return ", ".join(normalized_values) or None if isinstance(value, str): return value or None return value def markdown_table_cell_text(value: object) -> str: normalized = normalize_inline_whitespace(str(value or "")) if not normalized: return "" return normalized.replace("\\", "\\\\").replace("|", "\\|") MARKDOWN_EMAIL_ADDRESS_PATTERN = re.compile(r"^[^\s<>;,]+@[^\s<>;,]+$") def strip_wrapping_quotes(value: str) -> str: stripped = value.strip() if len(stripped) >= 2 and stripped[0] == stripped[-1] and stripped[0] in {"'", '"'}: return stripped[1:-1].strip() return stripped def split_markdown_person_email(value: object) -> tuple[str | None, str | None]: normalized = normalize_inline_whitespace(str(value or "")) if not normalized: return None, None if MARKDOWN_EMAIL_ADDRESS_PATTERN.fullmatch(normalized): return None, normalized bracket_match = re.fullmatch( r"(?P.+?)\s*<(?P[^\s<>;,]+@[^\s<>;,]+)>\s*", normalized, ) if bracket_match: name = strip_wrapping_quotes(normalize_inline_whitespace(bracket_match.group("name"))) email = normalize_inline_whitespace(bracket_match.group("email")) return (name or None), (email or None) trailing_match = re.fullmatch( r"(?P.+?)\s+(?P[^\s<>;,]+@[^\s<>;,]+)\s*", normalized, ) if trailing_match: name = strip_wrapping_quotes(normalize_inline_whitespace(trailing_match.group("name"))) email = normalize_inline_whitespace(trailing_match.group("email")) return (name or None), (email or None) return normalized, None def markdown_author_cell_text(value: object) -> str: name, email = split_markdown_person_email(value) if email and name and name != email: return f"{markdown_table_cell_text(name)}
    {markdown_table_cell_text(email)}" return markdown_table_cell_text(name or email) def markdown_link_text(value: str) -> str: return value.replace("\\", "\\\\").replace("[", "\\[").replace("]", "\\]") def format_search_markdown_date(value: object) -> str: parsed = parse_utc_timestamp(value) if parsed is None: return markdown_table_cell_text(value) return parsed.strftime("%Y-%m-%d %H:%M") def markdown_search_target_from_path(path_value: object) -> str | None: normalized = normalize_inline_whitespace(str(path_value or "")) if not normalized: return None if normalized.startswith("computer://"): return normalized return f"computer://{normalized}" def markdown_search_target(item: dict[str, object]) -> str | None: preview_targets = item.get("preview_targets") if isinstance(preview_targets, list): for preview_target in preview_targets: if not isinstance(preview_target, dict): continue if normalize_inline_whitespace(str(preview_target.get("label") or "")).lower() != "message": continue target = markdown_search_target_from_path( preview_target.get("abs_path") or preview_target.get("file_abs_path") ) if target: return target for key in ("preview_abs_path", "abs_path"): target = markdown_search_target_from_path(item.get(key)) if target: return target return None def summarize_child_content_type(item: dict[str, object], *, attachment: bool) -> str | None: if attachment: return "Unrecognized" child_kind = normalize_inline_whitespace(str(item.get("child_document_kind") or "")) if child_kind: return child_kind.replace("_", " ").title() return None def render_search_markdown_cell( item: dict[str, object], column_def: dict[str, str], *, child_prefix: str = "", child_content_type: str | None = None, standalone_child_parent: str | None = None, ) -> str: column_name = str(column_def["name"]) field_type = str(column_def["type"]) display_values = item.get("display_values") if isinstance(display_values, dict) and column_name in display_values and not child_prefix: value = display_values.get(column_name) else: value = summary_display_value(item, column_name, field_type) if column_name == "content_type" and child_content_type and not value: value = child_content_type if field_type == "date": return format_search_markdown_date(value) if column_name == "title": title_text = normalize_inline_whitespace(str(value or best_summary_title(item) or "Untitled")) if child_prefix: title_text = f"{child_prefix}{title_text}" if standalone_child_parent: title_text = f"{title_text} ({standalone_child_parent})" target = markdown_search_target(item) if target: return f"[{markdown_link_text(title_text)}]({target})" return markdown_table_cell_text(title_text) if column_name == "author": return markdown_author_cell_text(value) return markdown_table_cell_text(value) def render_search_markdown_row( item: dict[str, object], column_defs: list[dict[str, str]], *, child_prefix: str = "", child_content_type: str | None = None, standalone_child_parent: str | None = None, ) -> str: cells = [ render_search_markdown_cell( item, column_def, child_prefix=child_prefix, child_content_type=child_content_type, standalone_child_parent=standalone_child_parent, ) for column_def in column_defs ] return "| " + " | ".join(cells) + " |" SEARCH_HEADER_KEY_ORDER = ( "keyword", "bates", "filters", "custodians", "datasets", "from_run_id", "scope", "sort", "page", ) def compute_search_overview_line( results: object, *, browse_mode: str, ) -> str | None: if not isinstance(results, list) or not results: return None datasets: set[str] = set() custodians: set[str] = set() with_attachments = 0 flagged_empty = 0 for item in results: if not isinstance(item, dict): continue dataset_names_value = item.get("dataset_names") if isinstance(dataset_names_value, list): for name in dataset_names_value: if name is None: continue text = str(name).strip() if text: datasets.add(text) else: dataset_name_value = item.get("dataset_name") if dataset_name_value: for part in str(dataset_name_value).split(","): part = part.strip() if part: datasets.add(part) custodian_values = item.get("custodians") if isinstance(custodian_values, list) and custodian_values: for value in custodian_values: if value is None: continue text = str(value).strip() if text: custodians.add(text) else: custodian_value = item.get("custodian") if custodian_value: for part in str(custodian_value).split(","): part = part.strip() if part: custodians.add(part) attachments = item.get("attachments") attachment_count = 0 try: attachment_count = int(item.get("attachment_count") or 0) except (TypeError, ValueError): attachment_count = 0 if attachment_count > 0 or (isinstance(attachments, list) and attachments): with_attachments += 1 text_status = item.get("text_status") if text_status in {"empty", "no_text", "unavailable"}: flagged_empty += 1 parts: list[str] = [] if datasets: parts.append(f"{len(datasets)} dataset{'s' if len(datasets) != 1 else ''}") if custodians: parts.append(f"{len(custodians)} custodian{'s' if len(custodians) != 1 else ''}") if with_attachments: parts.append(f"{with_attachments} with attachments") if flagged_empty: parts.append(f"{flagged_empty} flagged (text_status=empty)") if not parts: return None label = browse_mode_result_label(browse_mode) parts.insert(0, f"{len(results)} {label} on this page") return "Overview: " + " \u00b7 ".join(parts) def build_search_footer_hints( payload: dict[str, object], results: object, *, page: int, total_pages: int, ) -> list[str]: hints: list[str] = [] nav: list[str] = [] if page < max(total_pages, 1): nav.append("`/retriever:next` for the next page") if page > 1: nav.append("`/retriever:previous` to go back") if nav: hints.append("Navigate: " + ", ".join(nav) + ".") if not isinstance(results, list) or not results: return hints dataset_counts: dict[str, int] = {} run_ids: set[str] = set() custodian_counts: dict[str, int] = {} for item in results: if not isinstance(item, dict): continue dataset_names_value = item.get("dataset_names") if isinstance(dataset_names_value, list): for name in dataset_names_value: if name is None: continue text = str(name).strip() if text: dataset_counts[text] = dataset_counts.get(text, 0) + 1 else: dataset_name_value = item.get("dataset_name") if dataset_name_value: for part in str(dataset_name_value).split(","): part = part.strip() if part: dataset_counts[part] = dataset_counts.get(part, 0) + 1 run_value = item.get("processing_run_id") or item.get("run_id") if run_value is not None and str(run_value).strip(): run_ids.add(str(run_value).strip()) custodian_values = item.get("custodians") if isinstance(custodian_values, list) and custodian_values: for value in custodian_values: if value is None: continue text = str(value).strip() if text: custodian_counts[text] = custodian_counts.get(text, 0) + 1 else: custodian_value = item.get("custodian") if custodian_value: for part in str(custodian_value).split(","): part = part.strip() if part: custodian_counts[part] = custodian_counts.get(part, 0) + 1 scope_value = payload.get("scope") active_scope = scope_value if isinstance(scope_value, dict) else {} has_dataset_scope = bool(coerce_scope_dataset_entries(active_scope.get("dataset"))) has_run_scope = active_scope.get("from_run_id") is not None existing_filter = str(active_scope.get("filter") or "") narrow: list[str] = [] if len(dataset_counts) > 1 and not has_dataset_scope: dominant_dataset = max(dataset_counts.items(), key=lambda item: item[1]) narrow.append( f"`/filter dataset_name = '{dominant_dataset[0]}'` to focus on one dataset" ) if len(run_ids) > 1 and not has_run_scope: narrow.append("`/from-run ` to focus on a specific processing run") if ( len(custodian_counts) > 1 and "custodian" not in existing_filter.lower() ): dominant_custodian = max(custodian_counts.items(), key=lambda item: item[1]) narrow.append( f"`/filter custodian = '{dominant_custodian[0]}'` to focus on one custodian" ) if narrow: hints.append("Narrow: " + "; ".join(narrow) + ".") return hints def render_search_markdown(payload: dict[str, object], column_defs: list[dict[str, str]]) -> str: lines: list[str] = [] browse_mode = normalize_browse_mode(payload.get("browse_mode")) header = payload.get("header") if isinstance(header, dict): for key in SEARCH_HEADER_KEY_ORDER: value = header.get(key) if payload_has_meaningful_value(value): lines.append(str(value)) for key, value in header.items(): if key in SEARCH_HEADER_KEY_ORDER: continue if payload_has_meaningful_value(value): lines.append(str(value)) overview_line = compute_search_overview_line( payload.get("results"), browse_mode=browse_mode, ) if overview_line: lines.append(overview_line) column_labels = [ markdown_table_cell_text(passive_field_label(column_def.get("name")) or column_def.get("name") or "") for column_def in column_defs ] lines.append("") lines.append("| " + " | ".join(column_labels) + " |") lines.append("|" + "|".join("---" for _ in column_labels) + "|") results = payload.get("results") if isinstance(results, list): for raw_item in results: if not isinstance(raw_item, dict): continue parent_context = None parent = raw_item.get("parent") if isinstance(parent, dict): parent_title = best_summary_title(parent) if parent_title: parent_context = f"parent: {parent_title}" child_kind = normalize_inline_whitespace(str(raw_item.get("child_document_kind") or "")) lines.append( render_search_markdown_row( raw_item, column_defs, child_prefix="↳ " if parent_context else "", child_content_type=summarize_child_content_type( raw_item, attachment=child_kind == CHILD_DOCUMENT_KIND_ATTACHMENT, ), standalone_child_parent=parent_context, ) ) attachments = raw_item.get("attachments") if isinstance(attachments, list): for attachment in attachments: if isinstance(attachment, dict): lines.append( render_search_markdown_row( attachment, column_defs, child_prefix="↳ ", child_content_type=summarize_child_content_type(attachment, attachment=True), ) ) child_documents = raw_item.get("child_documents") if isinstance(child_documents, list): for child in child_documents: if isinstance(child, dict): lines.append( render_search_markdown_row( child, column_defs, child_prefix="↳ ", child_content_type=summarize_child_content_type(child, attachment=False), ) ) total_hits = int(payload.get("total_hits") or 0) page = int(payload.get("page") or 1) per_page = int(payload.get("per_page") or DEFAULT_PAGE_SIZE) total_pages = int(payload.get("total_pages") or 1) start_index = 0 if total_hits == 0 else ((page - 1) * per_page) + 1 end_index = 0 if total_hits == 0 else min(total_hits, page * per_page) result_label = browse_mode_result_label(browse_mode, title_case=True) footer_lines = [f"{result_label} {start_index}\u2013{end_index} of {total_hits}."] footer_lines.extend( build_search_footer_hints( payload, payload.get("results"), page=page, total_pages=total_pages, ) ) lines.append("") lines.extend(footer_lines) return "\n".join(lines) def scope_dataset_name_suggestions(connection: sqlite3.Connection, dataset_name: str) -> list[str]: dataset_names = [summary["dataset_name"] for summary in list_dataset_summaries(connection)] return difflib.get_close_matches(normalize_inline_whitespace(dataset_name), dataset_names, n=3, cutoff=0.45) def resolve_scope_dataset_entries( connection: sqlite3.Connection, dataset_entries: object, ) -> list[dict[str, object]]: normalized_entries = coerce_scope_dataset_entries(dataset_entries) if len(normalized_entries) > MAX_SCOPE_DATASETS: raise RetrieverError(f"Scope datasets are capped at {MAX_SCOPE_DATASETS} entries.") resolved_entries: list[dict[str, object]] = [] seen_ids: set[int] = set() for entry in normalized_entries: dataset_id = int(entry["id"]) row = get_dataset_row_by_id(connection, dataset_id) if row is None: raise RetrieverError( f"Dataset id {dataset_id} ({entry['name']!r}) no longer exists. Clear with /dataset clear or replace with /dataset ." ) if dataset_id in seen_ids: continue seen_ids.add(dataset_id) resolved_entries.append({"id": dataset_id, "name": str(row["dataset_name"])}) return resolved_entries def resolve_scope_from_run_id(connection: sqlite3.Connection, from_run_id: object) -> int | None: if from_run_id in (None, ""): return None try: normalized_run_id = int(from_run_id) except (TypeError, ValueError) as exc: raise RetrieverError(f"Invalid run id: {from_run_id!r}") from exc row = connection.execute("SELECT id FROM runs WHERE id = ?", (normalized_run_id,)).fetchone() if row is None: raise RetrieverError( f"Run {normalized_run_id} referenced by scope.from_run_id no longer exists. Clear with /from-run clear." ) return normalized_run_id def build_scope_search_filters( connection: sqlite3.Connection, raw_scope: object, ) -> tuple[dict[str, object], list[str], list[object], list[object]]: scope = coerce_scope_payload(raw_scope) clauses = base_document_search_clauses() params: list[object] = [] filter_summary: list[object] = [] filter_expression = normalize_inline_whitespace(str(scope.get("filter") or "")) if filter_expression: clause, clause_params = compile_sql_filter_expression(connection, filter_expression) clauses.append(f"({clause})") params.extend(clause_params) filter_summary.append(filter_expression) dataset_entries = resolve_scope_dataset_entries(connection, scope.get("dataset")) if dataset_entries: placeholders = ", ".join("?" for _ in dataset_entries) clauses.append( "EXISTS (SELECT 1 FROM dataset_documents dd_scope " f"WHERE dd_scope.document_id = d.id AND dd_scope.dataset_id IN ({placeholders}))" ) params.extend(int(entry["id"]) for entry in dataset_entries) scope["dataset"] = dataset_entries from_run_id = resolve_scope_from_run_id(connection, scope.get("from_run_id")) if from_run_id is not None: clauses.append( "EXISTS (SELECT 1 FROM run_snapshot_documents rsd " "WHERE rsd.document_id = d.id AND rsd.run_id = ?)" ) params.append(from_run_id) scope["from_run_id"] = from_run_id return scope, clauses, params, filter_summary def stable_sort_results_by_field( results: list[dict[str, object]], field_name: str, order: str, ) -> list[dict[str, object]]: reverse = order.lower() == "desc" field_type = BUILTIN_FIELD_TYPES.get(field_name) non_null_items: list[tuple[object, dict[str, object]]] = [] null_items: list[dict[str, object]] = [] for item in results: row = item["row"] raw_value = row[field_name] if raw_value is None: null_items.append(item) continue normalized_value: object = raw_value if field_type == "date": parsed_value = parse_utc_timestamp(raw_value) if parsed_value is None: null_items.append(item) continue normalized_value = parsed_value elif isinstance(raw_value, str): normalized_value = raw_value.lower() non_null_items.append((normalized_value, item)) non_null_items.sort(key=lambda pair: pair[0], reverse=reverse) return [item for _, item in non_null_items] + null_items def coerce_sort_specs(raw_value: object) -> list[tuple[str, str]]: if not isinstance(raw_value, list): return [] normalized_specs: list[tuple[str, str]] = [] for item in raw_value: if not isinstance(item, (list, tuple)) or len(item) != 2: continue field_name = normalize_inline_whitespace(str(item[0] or "")) direction = normalize_inline_whitespace(str(item[1] or "")).lower() if not field_name or direction not in {"asc", "desc"}: continue normalized_specs.append((field_name, direction)) return normalized_specs def serialize_sort_specs(sort_specs: list[tuple[str, str]] | None) -> list[list[str]]: if not sort_specs: return [] return [[field_name, direction] for field_name, direction in sort_specs] def sort_specs_text(sort_specs: list[tuple[str, str]] | None) -> str | None: if not sort_specs: return None return ", ".join(f"{field_name} {direction}" for field_name, direction in sort_specs) def resolve_sort_field_name( connection: sqlite3.Connection, raw_field_name: str, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> str: field_def = resolve_browse_field_definition(connection, raw_field_name, browse_mode=browse_mode) if field_def.get("source") == "virtual": raise RetrieverError(f"Cannot sort by virtual filter field: {raw_field_name}") if str(field_def.get("sortable") or "true").lower() != "true": raise RetrieverError(f"Field '{field_def['field_name']}' cannot be used for sorting.") return str(field_def["field_name"]) def parse_slash_sort_specs( connection: sqlite3.Connection, raw_text: str, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> list[tuple[str, str]]: if not normalize_inline_whitespace(raw_text): raise RetrieverError("Usage: /sort ") parts = split_quoted_comma_values(raw_text) if not parts: raise RetrieverError("Usage: /sort ") sort_specs: list[tuple[str, str]] = [] for part in parts: tokens = shlex_split_slash_tail(part) if len(tokens) != 2: raise RetrieverError("Each sort entry must be exactly ' '.") field_name = resolve_sort_field_name(connection, tokens[0], browse_mode=browse_mode) direction = normalize_inline_whitespace(tokens[1]).lower() if direction not in {"asc", "desc"}: raise RetrieverError("Sort direction must be 'asc' or 'desc'.") sort_specs.append((field_name, direction)) return sort_specs def apply_sort_specs( results: list[dict[str, object]], sort_specs: list[tuple[str, str]] | None, ) -> list[dict[str, object]]: if not sort_specs: return results effective_specs = list(sort_specs) if not any(field_name == "id" for field_name, _ in effective_specs): effective_specs.append(("id", "asc")) sorted_results = list(results) for field_name, direction in reversed(effective_specs): if field_name == "id": sorted_results = sorted( sorted_results, key=lambda item: int(item["id"]), reverse=direction == "desc", ) continue sorted_results = stable_sort_results_by_field(sorted_results, field_name, direction) return sorted_results def stable_sort_summary_results_by_field( results: list[dict[str, object]], field_name: str, order: str, *, field_type: str, ) -> list[dict[str, object]]: reverse = order.lower() == "desc" non_null_items: list[tuple[object, dict[str, object]]] = [] null_items: list[dict[str, object]] = [] for item in results: raw_value = item.get(field_name) if raw_value in (None, ""): null_items.append(item) continue normalized_value: object = raw_value if field_type == "date": parsed_value = parse_utc_timestamp(raw_value) if parsed_value is None: null_items.append(item) continue normalized_value = parsed_value elif field_type == "integer": try: normalized_value = int(raw_value) except (TypeError, ValueError): null_items.append(item) continue elif field_type == "real": try: normalized_value = float(raw_value) except (TypeError, ValueError): null_items.append(item) continue elif isinstance(raw_value, str): normalized_value = raw_value.lower() non_null_items.append((normalized_value, item)) non_null_items.sort(key=lambda pair: pair[0], reverse=reverse) return [item for _, item in non_null_items] + null_items def apply_conversation_sort_specs( results: list[dict[str, object]], sort_specs: list[tuple[str, str]] | None, ) -> list[dict[str, object]]: if not sort_specs: return results effective_specs = list(sort_specs) if not any(field_name == "id" for field_name, _ in effective_specs): effective_specs.append(("id", "asc")) sorted_results = list(results) for field_name, direction in reversed(effective_specs): if field_name == "id": sorted_results = sorted( sorted_results, key=lambda item: int(item["id"]), reverse=direction == "desc", ) continue field_type = str(conversation_field_definition(field_name)["field_type"]) sorted_results = stable_sort_summary_results_by_field( sorted_results, field_name, direction, field_type=field_type, ) return sorted_results def resolve_scope_document_search( connection: sqlite3.Connection, raw_scope: object, *, sort_specs: list[tuple[str, str]] | None = None, ) -> dict[str, object]: scope, clauses, params, filter_summary = build_scope_search_filters(connection, raw_scope) keyword_query = normalize_inline_whitespace(str(scope.get("keyword") or "")) uses_relevance_scoring = keyword_query_uses_relevance_scoring(keyword_query) bates_scope = scope.get("bates") bates_query = format_scope_bates_value(bates_scope) bates_begin = normalize_inline_whitespace(str(bates_scope.get("begin") or "")) if isinstance(bates_scope, dict) else "" bates_end = normalize_inline_whitespace(str(bates_scope.get("end") or "")) if isinstance(bates_scope, dict) else "" bates_matches: dict[int, dict[str, object]] | None = None keyword_matches: dict[int, dict[str, object]] | None = None browse_matches: dict[int, dict[str, object]] | None = None if bates_query: bates_matches = search_bates(connection, bates_begin, bates_end, clauses, params) if keyword_query: keyword_matches = search_fts(connection, keyword_query, clauses, params) if bates_matches is None and keyword_matches is None: browse_matches = search_browse(connection, clauses, params) if bates_matches is not None and keyword_matches is not None: matches: dict[int, dict[str, object]] = {} for document_id, bates_match in bates_matches.items(): keyword_match = keyword_matches.get(document_id) if keyword_match is None: continue matches[document_id] = { "row": bates_match["row"], "rank": keyword_match.get("rank"), "snippet": keyword_match.get("snippet") or bates_match["snippet"], "bates_sort_key": bates_match.get("bates_sort_key"), } else: matches = bates_matches or keyword_matches or browse_matches or {} results = [ { "id": document_id, "rank": match["rank"], "snippet": match["snippet"], "bates_sort_key": match.get("bates_sort_key"), "row": match["row"], } for document_id, match in matches.items() ] stable_results = sorted(results, key=lambda item: int(item["id"])) if sort_specs: sorted_results = apply_sort_specs(stable_results, sort_specs) sort_name = sort_specs[0][0] order_name = sort_specs[0][1] sort_spec = sort_specs_text(sort_specs) or f"{sort_name} {order_name}" elif bates_query: sorted_results = sorted( stable_results, key=lambda item: item.get("bates_sort_key") or (1, "", 0, ""), ) sort_name = "bates" order_name = "asc" sort_spec = "bates asc" elif keyword_query and uses_relevance_scoring: sorted_results = stable_sort_results_by_field(stable_results, "date_created", "desc") sorted_results = sorted( sorted_results, key=lambda item: (item["rank"] is None, item["rank"]), ) sort_name = "relevance" order_name = "asc" sort_spec = "relevance asc" else: sorted_results = stable_sort_results_by_field(stable_results, "date_created", "desc") sort_name = "date_created" order_name = "desc" sort_spec = "date_created desc" query_label = keyword_query or bates_query or "" return { "scope": scope, "query": query_label, "filters": filter_summary, "sort": sort_name, "order": order_name, "sort_spec": sort_spec, "results": sorted_results, } def resolve_paged_scope_document_search( connection: sqlite3.Connection, raw_scope: object, *, sort_specs: list[tuple[str, str]] | None = None, offset: int, per_page: int, ) -> dict[str, object]: scope, clauses, params, filter_summary = build_scope_search_filters(connection, raw_scope) keyword_query = normalize_inline_whitespace(str(scope.get("keyword") or "")) uses_relevance_scoring = keyword_query_uses_relevance_scoring(keyword_query) bates_scope = scope.get("bates") bates_query = format_scope_bates_value(bates_scope) bates_begin = normalize_inline_whitespace(str(bates_scope.get("begin") or "")) if isinstance(bates_scope, dict) else "" bates_end = normalize_inline_whitespace(str(bates_scope.get("end") or "")) if isinstance(bates_scope, dict) else "" if sort_specs: sort_name = sort_specs[0][0] order_name = sort_specs[0][1] sort_spec = sort_specs_text(sort_specs) or f"{sort_name} {order_name}" elif bates_query: sort_name = "bates" order_name = "asc" sort_spec = "bates asc" elif keyword_query and uses_relevance_scoring: sort_name = "relevance" order_name = "asc" sort_spec = "relevance asc" else: sort_name = "date_created" order_name = "desc" sort_spec = "date_created desc" if bates_query and keyword_query: where_clause = " AND ".join(clauses) range_begin_expr = "COALESCE(d.begin_bates, d.control_number)" range_end_expr = "COALESCE(d.end_bates, d.control_number)" single_value = bates_begin == bates_end if single_value: match_clause = ( "d.control_number = ? OR d.begin_bates = ? OR d.end_bates = ? " f"OR ({range_begin_expr} <= ? AND {range_end_expr} >= ?)" ) match_params: list[object] = [bates_begin, bates_begin, bates_begin, bates_begin, bates_begin] else: match_clause = f"{range_begin_expr} <= ? AND {range_end_expr} >= ?" match_params = [bates_end, bates_begin] if sort_specs: order_by_sql = sql_order_by_for_sort_specs(connection, sort_specs, alias="d") else: order_by_sql = sql_bates_order_by(row_alias="bt", prioritize_rank=False, id_alias="d") cte_sql = f""" WITH chunk_matches AS ( SELECT d.id AS document_id, dc.text_content AS snippet_source, bm25(chunks_fts) AS rank, 0 AS source_priority FROM chunks_fts JOIN document_chunks dc ON dc.id = CAST(chunks_fts.chunk_id AS INTEGER) JOIN documents d ON d.id = dc.document_id WHERE chunks_fts MATCH ? AND {where_clause} ), metadata_matches AS ( SELECT d.id AS document_id, NULL AS snippet_source, bm25(documents_fts) AS rank, 1 AS source_priority FROM documents_fts JOIN documents d ON d.id = CAST(documents_fts.document_id AS INTEGER) WHERE documents_fts MATCH ? AND {where_clause} ), all_matches AS ( SELECT * FROM chunk_matches UNION ALL SELECT * FROM metadata_matches ), ranked_matches AS ( SELECT document_id, snippet_source, rank, source_priority, ROW_NUMBER() OVER ( PARTITION BY document_id ORDER BY rank ASC, source_priority ASC, document_id ASC ) AS row_number FROM all_matches ), best_matches AS ( SELECT document_id, snippet_source, rank FROM ranked_matches WHERE row_number = 1 ), bates_matches AS ( SELECT d.id AS document_id, {range_begin_expr} AS bates_sort_value FROM documents d WHERE {where_clause} AND ({match_clause}) ) """ def combined_params(fts_query: str) -> list[object]: return [fts_query, *params, fts_query, *params, *params, *match_params] effective_query = keyword_query count_sql = f""" {cte_sql} SELECT COUNT(*) AS total_hits FROM best_matches bm JOIN bates_matches bt ON bt.document_id = bm.document_id """ try: count_row = connection.execute(count_sql, combined_params(effective_query)).fetchone() except sqlite3.OperationalError: effective_query = f'"{keyword_query}"' count_row = connection.execute(count_sql, combined_params(effective_query)).fetchone() total_hits = int(count_row["total_hits"] or 0) if count_row is not None else 0 rows = connection.execute( f""" {cte_sql} SELECT d.*, bm.rank AS rank, bm.snippet_source AS snippet_source, bt.bates_sort_value AS bates_sort_value FROM best_matches bm JOIN bates_matches bt ON bt.document_id = bm.document_id JOIN documents d ON d.id = bm.document_id ORDER BY {order_by_sql} LIMIT ? OFFSET ? """, [*combined_params(effective_query), per_page, offset], ).fetchall() selection_page = { "total_hits": total_hits, "results": [ { "id": int(row["id"]), "rank": float(row["rank"]) if row["rank"] is not None else None, "snippet": make_snippet( row["snippet_source"] if row["snippet_source"] else metadata_snippet(row), keyword_query, ), "bates_sort_key": bates_sort_key(row["bates_sort_value"] or row["control_number"]), "row": row, } for row in rows ], } elif bates_query: if sort_specs: selection_page = search_bates_page( connection, bates_begin, bates_end, clauses, params, limit=per_page, offset=offset, order_by_sql=sql_order_by_for_sort_specs(connection, sort_specs, alias="bm"), ) else: selection_page = search_bates_page( connection, bates_begin, bates_end, clauses, params, limit=per_page, offset=offset, order_by_sql=sql_bates_order_by(row_alias="bm", prioritize_rank=False), ) elif keyword_query: if sort_specs: order_by_sql = sql_order_by_for_sort_specs(connection, sort_specs, alias="d") else: if uses_relevance_scoring: order_by_sql = sql_relevance_order_by(connection, row_alias="d", rank_expr="bm.rank") else: order_by_sql = sql_order_by_for_sort_specs(connection, [("date_created", "desc")], alias="d") selection_page = search_fts_page( connection, keyword_query, clauses, params, limit=per_page, offset=offset, order_by_sql=order_by_sql, ) else: if sort_specs: order_by_sql = sql_order_by_for_sort_specs(connection, sort_specs, alias="d") else: order_by_sql = sql_order_by_for_sort_specs(connection, [("date_created", "desc")], alias="d") selection_page = search_browse_page( connection, clauses, params, limit=per_page, offset=offset, order_by_sql=order_by_sql, ) query_label = keyword_query or bates_query or "" return { "scope": scope, "query": query_label, "filters": filter_summary, "sort": sort_name, "order": order_name, "sort_spec": sort_spec, "results": selection_page["results"], "total_hits": int(selection_page["total_hits"]), } def resolve_paged_scope_conversation_search( connection: sqlite3.Connection, paths: dict[str, Path], raw_scope: object, *, sort_specs: list[tuple[str, str]] | None = None, offset: int, per_page: int, ) -> dict[str, object]: selection = resolve_scope_document_search(connection, raw_scope) scope = coerce_scope_payload(selection.get("scope")) keyword_query = normalize_inline_whitespace(str(scope.get("keyword") or "")) uses_relevance_scoring = keyword_query_uses_relevance_scoring(keyword_query) bates_query = format_scope_bates_value(scope.get("bates")) grouped_matches: dict[int, dict[str, object]] = {} for match in selection["results"]: row = match["row"] if row["conversation_id"] is None: continue if normalize_inline_whitespace(str(row["child_document_kind"] or "")) == CHILD_DOCUMENT_KIND_ATTACHMENT: continue conversation_id = int(row["conversation_id"]) payload = grouped_matches.setdefault( conversation_id, { "matching_document_ids": set(), "snippet": "", "rank": None, "bates_sort_key": None, }, ) matching_document_ids = payload["matching_document_ids"] assert isinstance(matching_document_ids, set) matching_document_ids.add(int(match["id"])) snippet = normalize_inline_whitespace(str(match.get("snippet") or "")) rank = match.get("rank") if rank is not None and (payload["rank"] is None or rank < payload["rank"]): payload["rank"] = rank if snippet: payload["snippet"] = snippet elif not payload["snippet"] and snippet: payload["snippet"] = snippet bates_sort_key = match.get("bates_sort_key") if bates_sort_key is not None and ( payload["bates_sort_key"] is None or bates_sort_key < payload["bates_sort_key"] ): payload["bates_sort_key"] = bates_sort_key if not payload["snippet"] and snippet: payload["snippet"] = snippet conversation_ids = sorted(grouped_matches) if not conversation_ids: return { "scope": scope, "query": selection["query"], "filters": selection["filters"], "sort": "last_activity", "order": "desc", "sort_spec": "last_activity desc", "results": [], "total_hits": 0, } placeholders = ", ".join("?" for _ in conversation_ids) conversation_rows = connection.execute( f""" SELECT id, source_kind, conversation_type, display_name FROM conversations WHERE id IN ({placeholders}) """, conversation_ids, ).fetchall() conversation_rows_by_id = {int(row["id"]): row for row in conversation_rows} summary_documents_by_conversation = load_conversation_summary_documents(connection, conversation_ids) summary_results: list[dict[str, object]] = [] for conversation_id in conversation_ids: conversation_row = conversation_rows_by_id.get(conversation_id) documents = summary_documents_by_conversation.get(conversation_id, []) if conversation_row is None or not documents: continue first_activity, last_activity = conversation_preview_bounds(documents) matching_document_ids = grouped_matches[conversation_id]["matching_document_ids"] assert isinstance(matching_document_ids, set) summary_results.append( { "id": conversation_id, "conversation_id": conversation_id, "conversation_type": normalize_inline_whitespace(str(conversation_row["conversation_type"] or "")), "title": normalize_inline_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {conversation_id}", "participants": conversation_preview_participants(documents), "first_activity": first_activity, "last_activity": last_activity, "document_count": len(documents), "matching_document_count": len(matching_document_ids), "source_kind": normalize_inline_whitespace(str(conversation_row["source_kind"] or "")) or None, "snippet": str(grouped_matches[conversation_id]["snippet"] or ""), "rank": grouped_matches[conversation_id]["rank"], "_bates_sort_key": grouped_matches[conversation_id]["bates_sort_key"], "_document_ids": [int(document["id"]) for document in documents], } ) stable_results = sorted(summary_results, key=lambda item: int(item["id"])) if sort_specs: sorted_results = apply_conversation_sort_specs(stable_results, sort_specs) sort_name = sort_specs[0][0] order_name = sort_specs[0][1] sort_spec = sort_specs_text(sort_specs) or f"{sort_name} {order_name}" elif bates_query: sorted_results = sorted( stable_results, key=lambda item: item.get("_bates_sort_key") or (1, "", 0, ""), ) sort_name = "bates" order_name = "asc" sort_spec = "bates asc" elif keyword_query and uses_relevance_scoring: sorted_results = stable_sort_summary_results_by_field( stable_results, "last_activity", "desc", field_type="date", ) sorted_results = stable_sort_summary_results_by_field( sorted_results, "rank", "asc", field_type="real", ) sort_name = "relevance" order_name = "asc" sort_spec = "relevance asc" else: sorted_results = stable_sort_summary_results_by_field( stable_results, "last_activity", "desc", field_type="date", ) sort_name = "last_activity" order_name = "desc" sort_spec = "last_activity desc" total_hits = len(sorted_results) paged_results = [dict(item) for item in sorted_results[offset: offset + per_page]] paged_document_ids = [ document_id for item in paged_results for document_id in item.get("_document_ids", []) if isinstance(document_id, int) ] dataset_memberships = fetch_dataset_memberships_for_document_ids(connection, paged_document_ids) for item in paged_results: dataset_ids: list[int] = [] dataset_names: list[str] = [] for document_id in item.pop("_document_ids", []): memberships = dataset_memberships.get(int(document_id), {"ids": [], "names": []}) for dataset_id in memberships["ids"]: normalized_dataset_id = int(dataset_id) if normalized_dataset_id not in dataset_ids: dataset_ids.append(normalized_dataset_id) for dataset_name in memberships["names"]: normalized_dataset_name = str(dataset_name) if normalized_dataset_name not in dataset_names: dataset_names.append(normalized_dataset_name) item["dataset_ids"] = dataset_ids item["dataset_names"] = dataset_names if len(dataset_names) == 1: item["dataset_name"] = dataset_names[0] item["metadata"] = { "conversation_type": item.get("conversation_type"), "title": item.get("title"), "participants": item.get("participants"), "first_activity": item.get("first_activity"), "last_activity": item.get("last_activity"), "document_count": item.get("document_count"), "source_kind": item.get("source_kind"), } item.update(conversation_path_payload(paths, connection, int(item["id"]))) item.pop("_bates_sort_key", None) return { "scope": scope, "query": selection["query"], "filters": selection["filters"], "sort": sort_name, "order": order_name, "sort_spec": sort_spec, "results": paged_results, "total_hits": total_hits, } def search_with_scope( root: Path, raw_scope: object, *, page: int = 1, per_page: int = DEFAULT_PAGE_SIZE, offset: int | None = None, sort_specs: list[tuple[str, str]] | None = None, display_column_defs: list[dict[str, str]] | None = None, warnings: list[str] | None = None, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> dict[str, object]: if page < 1: raise RetrieverError("Page must be >= 1.") if per_page < 1: raise RetrieverError("per-page must be >= 1.") per_page = min(per_page, MAX_PAGE_SIZE) if offset is not None and offset < 0: raise RetrieverError("Offset must be >= 0.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) effective_browse_mode = normalize_browse_mode(browse_mode) requested_offset = (page - 1) * per_page if offset is None else offset effective_display_column_defs = display_column_defs or default_display_column_definitions( connection, browse_mode=effective_browse_mode, ) if effective_browse_mode == BROWSE_MODE_CONVERSATIONS: selection = resolve_paged_scope_conversation_search( connection, paths, raw_scope, sort_specs=sort_specs, offset=requested_offset, per_page=per_page, ) total_hits = int(selection["total_hits"]) total_pages = max(1, (total_hits + per_page - 1) // per_page) start = requested_offset if total_hits > 0 and start >= total_hits: start = (total_pages - 1) * per_page selection = resolve_paged_scope_conversation_search( connection, paths, raw_scope, sort_specs=sort_specs, offset=start, per_page=per_page, ) start = max(0, start) page = (start // per_page) + 1 paged_results = [dict(item) for item in selection["results"]] for item in paged_results: item["display_values"] = build_summary_display_values(item, effective_display_column_defs) payload = { "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], "sort_spec": selection["sort_spec"], "browse_mode": effective_browse_mode, "page": page, "per_page": per_page, "offset": start, "total_hits": total_hits, "total_pages": total_pages, "results": paged_results, "scope": selection["scope"], "display": build_display_payload(effective_display_column_defs, per_page), } payload["header"] = build_search_header_payload(selection["scope"], payload) if warnings: payload["warnings"] = warnings payload["rendered_markdown"] = render_search_markdown(payload, effective_display_column_defs) return payload results: list[dict[str, object]] = [] selection = resolve_paged_scope_document_search( connection, raw_scope, sort_specs=sort_specs, offset=requested_offset, per_page=per_page, ) for match in selection["results"]: row = match["row"] results.append( { "id": int(match["id"]), "control_number": row["control_number"], "dataset_id": row["dataset_id"], "parent_document_id": row["parent_document_id"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "source_folder_path": row["source_folder_path"], "production_id": row["production_id"], **document_path_payload(paths, connection, row), "file_name": row["file_name"], "file_type": row["file_type"], "snippet": str(match["snippet"]), "rank": match["rank"], "bates_sort_key": match.get("bates_sort_key"), "metadata": { "author": row["author"], "begin_attachment": row["begin_attachment"], "begin_bates": row["begin_bates"], "content_type": row["content_type"], "custodian": document_custodian_display_text_from_row(row), "custodians": document_custodian_values_from_row(row), "dataset_id": row["dataset_id"], "date_created": row["date_created"], "date_modified": row["date_modified"], "end_attachment": row["end_attachment"], "end_bates": row["end_bates"], "page_count": row["page_count"], "participants": row["participants"], "recipients": row["recipients"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "source_folder_path": row["source_folder_path"], "subject": row["subject"], "title": row["title"], "updated_at": row["updated_at"], }, "manual_field_locks": normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN]), "row": row, } ) total_hits = int(selection["total_hits"]) total_pages = max(1, (total_hits + per_page - 1) // per_page) start = requested_offset if total_hits > 0 and start >= total_hits: start = (total_pages - 1) * per_page selection = resolve_paged_scope_document_search( connection, raw_scope, sort_specs=sort_specs, offset=start, per_page=per_page, ) results = [] for match in selection["results"]: row = match["row"] results.append( { "id": int(match["id"]), "control_number": row["control_number"], "dataset_id": row["dataset_id"], "parent_document_id": row["parent_document_id"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "source_folder_path": row["source_folder_path"], "production_id": row["production_id"], **document_path_payload(paths, connection, row), "file_name": row["file_name"], "file_type": row["file_type"], "snippet": str(match["snippet"]), "rank": match["rank"], "bates_sort_key": match.get("bates_sort_key"), "metadata": { "author": row["author"], "begin_attachment": row["begin_attachment"], "begin_bates": row["begin_bates"], "content_type": row["content_type"], "custodian": document_custodian_display_text_from_row(row), "custodians": document_custodian_values_from_row(row), "dataset_id": row["dataset_id"], "date_created": row["date_created"], "date_modified": row["date_modified"], "end_attachment": row["end_attachment"], "end_bates": row["end_bates"], "page_count": row["page_count"], "participants": row["participants"], "recipients": row["recipients"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "source_folder_path": row["source_folder_path"], "subject": row["subject"], "title": row["title"], "updated_at": row["updated_at"], }, "manual_field_locks": normalize_string_list(row[MANUAL_FIELD_LOCKS_COLUMN]), "row": row, } ) start = max(0, start) page = (start // per_page) + 1 paged_results = results paged_rows = [item["row"] for item in paged_results] production_names = fetch_production_names(connection, paged_rows) dataset_memberships = fetch_document_dataset_memberships(connection, paged_rows) attachment_summaries = fetch_attachment_summaries( connection, paths, [int(row["id"]) for row in paged_rows if row["parent_document_id"] is None], ) parent_summaries = fetch_parent_summaries( connection, [row for row in paged_rows if row["parent_document_id"] is not None], ) for item in paged_results: row = item["row"] custodian_values = document_custodian_values_from_row(row) custodian_text = ", ".join(custodian_values) if custodian_values else None item["custodian"] = custodian_text item["custodians"] = custodian_values item["metadata"]["custodian"] = custodian_text item["metadata"]["custodians"] = custodian_values memberships = dataset_memberships.get(int(row["id"]), {"ids": [], "names": []}) dataset_ids = [int(dataset_id) for dataset_id in memberships["ids"]] dataset_names = [str(dataset_name) for dataset_name in memberships["names"]] item["dataset_ids"] = dataset_ids item["dataset_names"] = dataset_names item["metadata"]["dataset_ids"] = dataset_ids item["metadata"]["dataset_names"] = dataset_names if len(dataset_ids) == 1: item["dataset_id"] = dataset_ids[0] item["metadata"]["dataset_id"] = dataset_ids[0] else: item["dataset_id"] = None item["metadata"]["dataset_id"] = None if len(dataset_names) == 1: item["dataset_name"] = dataset_names[0] item["metadata"]["dataset_name"] = dataset_names[0] if row["production_id"] is not None: item["production_name"] = production_names.get(int(row["production_id"])) item["metadata"]["production_name"] = production_names.get(int(row["production_id"])) if row["parent_document_id"] is None: attachments = attachment_summaries.get(int(row["id"]), []) item["attachment_count"] = len(attachments) item["attachments"] = attachments else: item["parent"] = parent_summaries.get(int(row["parent_document_id"])) item["display_values"] = build_search_result_display_values(row, item, effective_display_column_defs) item.pop("bates_sort_key", None) item.pop("row", None) payload = { "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], "sort_spec": selection["sort_spec"], "browse_mode": effective_browse_mode, "page": page, "per_page": per_page, "offset": start, "total_hits": total_hits, "total_pages": total_pages, "results": paged_results, "scope": selection["scope"], "display": build_display_payload(effective_display_column_defs, per_page), } payload["header"] = build_search_header_payload(selection["scope"], payload) if warnings: payload["warnings"] = warnings payload["rendered_markdown"] = render_search_markdown(payload, effective_display_column_defs) return payload finally: connection.close() def session_page_size( session_state: dict[str, object], *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> int: display = session_display_state(session_state, browse_mode=browse_mode) if not display: return DEFAULT_PAGE_SIZE page_size = display.get("page_size") if not isinstance(page_size, int) or page_size < 1: return DEFAULT_PAGE_SIZE return min(page_size, MAX_PAGE_SIZE) def session_browsing_state( session_state: dict[str, object], *, browse_mode: str | None = None, ) -> dict[str, object]: effective_browse_mode = normalize_browse_mode(browse_mode or session_browse_mode(session_state)) browsing = session_state.get("browsing") if not isinstance(browsing, dict): return {} branch = browsing.get(effective_browse_mode) return branch if isinstance(branch, dict) else {} def session_sort_specs( session_state: dict[str, object], *, browse_mode: str | None = None, ) -> list[tuple[str, str]]: return coerce_sort_specs(session_browsing_state(session_state, browse_mode=browse_mode).get("sort")) def saved_scope_summaries(paths: dict[str, Path]) -> list[dict[str, object]]: saved_scopes_state = read_saved_scopes_state(paths) scopes = saved_scopes_state.get("scopes") if not isinstance(scopes, dict): return [] return [ { "name": str(scope_name), "scope": coerce_saved_scope_payload(scope_payload), } for scope_name, scope_payload in sorted( scopes.items(), key=lambda item: (normalize_saved_scope_name(str(item[0])), str(item[0])), ) ] def sortable_field_entries( connection: sqlite3.Connection, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> list[dict[str, object]]: seen_names: set[str] = set() entries: list[dict[str, object]] = [] for raw_field_name in known_browse_field_names(connection, browse_mode=browse_mode): field_def = resolve_browse_field_definition(connection, raw_field_name, browse_mode=browse_mode) if field_def.get("source") == "virtual": continue if str(field_def.get("sortable") or "true").lower() != "true": continue canonical_name = str(field_def["field_name"]) if canonical_name in seen_names: continue seen_names.add(canonical_name) entries.append( catalog_field_entry( canonical_name, str(field_def["field_type"]), source=str(field_def["source"]), instruction=field_def.get("instruction"), displayable=str(field_def.get("displayable") or "").lower() == "true", ) ) return sorted(entries, key=lambda item: (str(item["name"]).lower(), str(item["name"]))) def displayable_field_entries( connection: sqlite3.Connection, *, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> list[dict[str, object]]: entries: list[dict[str, object]] = [] for field_name in displayable_field_names(connection, browse_mode=browse_mode): field_def = resolve_browse_field_definition(connection, field_name, browse_mode=browse_mode) entries.append( catalog_field_entry( str(field_def["field_name"]), str(field_def["field_type"]), source=str(field_def["source"]), instruction=field_def.get("instruction"), displayable=True, ) ) return entries def active_sort_payload( scope: dict[str, object], session_state: dict[str, object], *, browse_mode: str | None = None, ) -> dict[str, object]: effective_browse_mode = normalize_browse_mode(browse_mode or session_browse_mode(session_state)) sort_specs = session_sort_specs(session_state, browse_mode=effective_browse_mode) if sort_specs: field_name, direction = sort_specs[0] return { "sort": field_name, "order": direction, "sort_spec": sort_specs_text(sort_specs) or f"{field_name} {direction}", "sort_source": "override", "sort_override": serialize_sort_specs(sort_specs), } if isinstance(scope.get("bates"), dict): return { "sort": "bates", "order": "asc", "sort_spec": "bates asc", "sort_source": "default", } keyword_query = normalize_inline_whitespace(str(scope.get("keyword") or "")) if keyword_query and keyword_query_uses_relevance_scoring(keyword_query): return { "sort": "relevance", "order": "asc", "sort_spec": "relevance asc", "sort_source": "default", } if effective_browse_mode == BROWSE_MODE_CONVERSATIONS: return { "sort": "last_activity", "order": "desc", "sort_spec": "last_activity desc", "sort_source": "default", } if effective_browse_mode == BROWSE_MODE_ENTITIES: return { "sort": "document_count", "order": "desc", "sort_spec": "document_count desc, label asc, id asc", "sort_source": "default", } return { "sort": "date_created", "order": "desc", "sort_spec": "date_created desc", "sort_source": "default", } def active_page_payload( session_state: dict[str, object], *, browse_mode: str | None = None, ) -> dict[str, int | str]: effective_browse_mode = normalize_browse_mode(browse_mode or session_browse_mode(session_state)) per_page = session_page_size(session_state, browse_mode=effective_browse_mode) browsing = session_browsing_state(session_state, browse_mode=effective_browse_mode) offset = int(browsing.get("offset") or 0) total_known = int(browsing.get("total_known") or 0) total_pages = max(1, (total_known + per_page - 1) // per_page) if total_known > 0 and offset >= total_known: offset = max(0, (total_pages - 1) * per_page) return { "browse_mode": effective_browse_mode, "page": (offset // per_page) + 1, "per_page": per_page, "offset": offset, "total_known": total_known, "total_pages": total_pages, } def parse_slash_command_text(raw_command: str) -> tuple[str, str]: command_text = raw_command.strip() if not command_text: raise RetrieverError("Slash command cannot be empty.") if not command_text.startswith("/"): raise RetrieverError("Slash commands must begin with '/'.") command_body = command_text[1:] command_name, _, tail = command_body.partition(" ") return command_name, tail.lstrip() def bates_scope_text(scope: dict[str, object]) -> str | None: bates = scope.get("bates") if not isinstance(bates, dict): return None begin = normalize_inline_whitespace(str(bates.get("begin") or "")) end = normalize_inline_whitespace(str(bates.get("end") or "")) if not begin or not end: return None return f"{begin}-{end}" def summarize_scope_inline(scope: dict[str, object]) -> str: parts: list[str] = [] keyword = normalize_inline_whitespace(str(scope.get("keyword") or "")) if keyword: parts.append(f"keyword={keyword}") bates_text = bates_scope_text(scope) if bates_text: parts.append(f"bates={bates_text}") filter_expression = normalize_inline_whitespace(str(scope.get("filter") or "")) if filter_expression: parts.append(f"filter={filter_expression}") dataset_entries = coerce_scope_dataset_entries(scope.get("dataset")) if dataset_entries: parts.append("dataset=" + ", ".join(str(entry["name"]) for entry in dataset_entries)) from_run_id = scope.get("from_run_id") if from_run_id is not None: parts.append(f"from-run={from_run_id}") return "; ".join(parts) if parts else "(none)" def format_dataset_size_summary(item: dict[str, object]) -> str: size_bytes = item.get("size_bytes") if size_bytes is None: return "—" try: size_value = int(size_bytes) except (TypeError, ValueError): return "—" unit_index = 0 display_value = float(size_value) units = ["B", "KB", "MB", "GB", "TB"] while display_value >= 1024.0 and unit_index < len(units) - 1: display_value /= 1024.0 unit_index += 1 if unit_index == 0: size_text = f"{int(display_value)} {units[unit_index]}" else: precision = 1 if display_value < 10 else 0 size_text = f"{display_value:.{precision}f}".rstrip("0").rstrip(".") + f" {units[unit_index]}" size_basis = normalize_inline_whitespace(str(item.get("size_basis") or "")).lower() if size_basis == "container": return size_text document_count = int(item.get("document_count") or 0) sized_document_count = int(item.get("sized_document_count") or 0) if document_count > 0 and 0 < sized_document_count < document_count: return f"{size_text} ({sized_document_count}/{document_count} sized)" return size_text def format_dataset_custodian_summary(item: dict[str, object]) -> str: raw_values = item.get("custodians") if not isinstance(raw_values, list): return "—" values = [ normalize_inline_whitespace(str(value or "")) for value in raw_values if normalize_inline_whitespace(str(value or "")) ] if not values: return "—" if len(values) <= 2: return ", ".join(values) return ", ".join(values[:2]) + f" +{len(values) - 2}" def escape_markdown_table_cell(value: object) -> str: text = normalize_inline_whitespace(str(value or "")) if not text: return "—" return text.replace("\\", "\\\\").replace("|", "\\|") def browse_mode_result_label(browse_mode: str, *, title_case: bool = False) -> str: effective_browse_mode = normalize_browse_mode(browse_mode) if effective_browse_mode == BROWSE_MODE_CONVERSATIONS: return "Conversations" if title_case else "conversations" if effective_browse_mode == BROWSE_MODE_ENTITIES: return "Entities" if title_case else "entities" return "Documents" if title_case else "docs" def render_slash_read_only_output(raw_command: str, payload: dict[str, object]) -> str | None: command_name, normalized_tail = parse_slash_command_text(raw_command) browse_mode = normalize_browse_mode(payload.get("browse_mode")) result_label = browse_mode_result_label(browse_mode) if command_name == "field": field_args = shlex_split_slash_tail(normalized_tail) if normalized_tail else [] if not field_args or field_args == ["list"]: fields = payload.get("fields") if not isinstance(fields, list) or not fields: return "Custom fields: (none)" lines = ["Custom fields:"] for item in fields: if not isinstance(item, dict): continue field_name = normalize_inline_whitespace(str(item.get("field_name") or "")) field_type = normalize_inline_whitespace(str(item.get("field_type") or "")) if not field_name or not field_type: continue document_count = int(item.get("documents_with_values") or 0) instruction = normalize_inline_whitespace(str(item.get("instruction") or "")) suffix = f": {instruction}" if instruction else "" lines.append(f"- {field_name} ({field_type}, {document_count} docs){suffix}") return "\n".join(lines) return None if command_name == "fill" and ( payload.get("status") == "confirm_required" or bool(payload.get("dry_run")) ): field_name = normalize_inline_whitespace(str(payload.get("field_name") or "")) if not field_name: return None action = "clear" if bool(payload.get("clear")) else f"fill {field_name}={payload.get('value')!r}" document_count = int(payload.get("would_write") or 0) summary = f"Preview: {action} on {document_count} document" if document_count != 1: summary += "s" if payload.get("status") == "confirm_required": summary += ". Re-run with --confirm to apply." return summary if command_name in {"next", "previous", "documents", "conversations", "entities"}: rendered_markdown = payload.get("rendered_markdown") if payload_has_meaningful_value(rendered_markdown): return str(rendered_markdown) if command_name == "search" and not normalized_tail: keyword = normalize_inline_whitespace(str(payload.get("keyword") or "")) return f"Search: {keyword}" if keyword else "Search: (none)" if command_name == "filter" and not normalized_tail: filter_expression = normalize_inline_whitespace(str(payload.get("filter") or "")) return f"Filter: {filter_expression}" if filter_expression else "Filter: (none)" if command_name == "bates" and not normalized_tail: bates = payload.get("bates") if isinstance(bates, dict): begin = normalize_inline_whitespace(str(bates.get("begin") or "")) end = normalize_inline_whitespace(str(bates.get("end") or "")) if begin and end: return f"Bates: {begin}-{end}" return "Bates: (none)" if command_name == "from-run" and not normalized_tail: from_run_id = payload.get("from_run_id") return f"From run: {from_run_id}" if from_run_id is not None else "From run: (none)" if command_name == "scope": scope_args = shlex_split_slash_tail(normalized_tail) if normalized_tail else [] if not scope_args: scope = coerce_scope_payload(payload.get("scope")) lines = ["Scope:"] if not scope: lines[0] = "Scope: (none)" else: keyword = normalize_inline_whitespace(str(scope.get("keyword") or "")) if keyword: lines.append(f"- keyword: {keyword}") bates_text = bates_scope_text(scope) if bates_text: lines.append(f"- bates: {bates_text}") filter_expression = normalize_inline_whitespace(str(scope.get("filter") or "")) if filter_expression: lines.append(f"- filter: {filter_expression}") dataset_entries = coerce_scope_dataset_entries(scope.get("dataset")) if dataset_entries: lines.append("- dataset: " + ", ".join(str(entry["name"]) for entry in dataset_entries)) from_run_id = scope.get("from_run_id") if from_run_id is not None: lines.append(f"- from-run: {from_run_id}") return "\n".join(lines) if scope_args == ["list"]: saved_scopes = payload.get("saved_scopes") if not isinstance(saved_scopes, list) or not saved_scopes: return "Saved scopes: (none)" lines = ["Saved scopes:"] for item in saved_scopes: if not isinstance(item, dict): continue name = normalize_inline_whitespace(str(item.get("name") or "")) if not name: continue scope = coerce_scope_payload(item.get("scope")) lines.append(f"- {name}: {summarize_scope_inline(scope)}") return "\n".join(lines) return None if command_name == "dataset": dataset_args = shlex_split_slash_tail(normalized_tail) if normalized_tail else [] if not dataset_args: dataset_entries = coerce_scope_dataset_entries(payload.get("dataset")) if not dataset_entries: return "Dataset: (none)" return "Dataset: " + ", ".join(str(entry["name"]) for entry in dataset_entries) if dataset_args == ["list"]: datasets = payload.get("datasets") if not isinstance(datasets, list) or not datasets: return "Datasets: (none)" lines = [ "Datasets:", "", "| Dataset | Docs | Size | Custodians |", "| --- | ---: | --- | --- |", ] for item in datasets: if not isinstance(item, dict): continue dataset_name = normalize_inline_whitespace(str(item.get("dataset_name") or "")) if not dataset_name: continue document_count = str(int(item.get("document_count") or 0)) lines.append( "| " + " | ".join( [ escape_markdown_table_cell(dataset_name), escape_markdown_table_cell(document_count), escape_markdown_table_cell(format_dataset_size_summary(item)), escape_markdown_table_cell(format_dataset_custodian_summary(item)), ] ) + " |" ) return "\n".join(lines) return None if command_name == "sort": if not normalized_tail: sort_spec = normalize_inline_whitespace(str(payload.get("sort_spec") or "")) sort_source = normalize_inline_whitespace(str(payload.get("sort_source") or "")) or "default" return f"Sort: {sort_spec} ({sort_source})" if sort_spec else "Sort: (none)" if normalized_tail == "list": sortable_fields = payload.get("sortable_fields") if not isinstance(sortable_fields, list) or not sortable_fields: return "Sortable fields: (none)" lines = ["Sortable fields:"] for item in sortable_fields: if not isinstance(item, dict): continue field_name = normalize_inline_whitespace(str(item.get("name") or "")) if field_name: lines.append(f"- {field_name}") return "\n".join(lines) return None if command_name == "columns": if not normalized_tail: display = payload.get("display") if isinstance(payload.get("display"), dict) else {} columns = display.get("columns") if isinstance(display.get("columns"), list) else [] page_size = display.get("page_size") lines = [ "Columns: " + ", ".join(str(column) for column in columns) if columns else "Columns: (none)" ] if page_size is not None: lines.append(f"Page size: {page_size}") warnings = payload.get("warnings") if isinstance(warnings, list): for warning in warnings: warning_text = normalize_inline_whitespace(str(warning or "")) if warning_text: lines.append(f"Warning: {warning_text}") return "\n".join(lines) if normalized_tail == "list": columns = payload.get("columns") if not isinstance(columns, list) or not columns: return "Displayable columns: (none)" lines = ["Displayable columns:"] for item in columns: if not isinstance(item, dict): continue field_name = normalize_inline_whitespace(str(item.get("name") or "")) if field_name: lines.append(f"- {field_name}") return "\n".join(lines) return None if command_name == "page-size" and not normalized_tail: return f"Page size: {int(payload.get('page_size') or 0)}" if command_name == "page" and not normalized_tail: page = int(payload.get("page") or 1) per_page = int(payload.get("per_page") or 0) offset = int(payload.get("offset") or 0) total_known = int(payload.get("total_known") or 0) total_pages = int(payload.get("total_pages") or 1) if total_known > 0: first_doc = offset + 1 last_doc = min(offset + per_page, total_known) else: first_doc = 0 last_doc = 0 return f"Page: {page} of {total_pages} ({result_label} {first_doc}-{last_doc} of {total_known})" return None def render_list_fields_table(payload: dict[str, object]) -> str: fields = payload.get("fields") if not isinstance(fields, list) or not fields: return "Custom fields: (none)" lines = ["NAME | TYPE | DOCS | DESCRIPTION"] for item in fields: if not isinstance(item, dict): continue field_name = normalize_inline_whitespace(str(item.get("field_name") or "")) field_type = normalize_inline_whitespace(str(item.get("field_type") or "")) document_count = int(item.get("documents_with_values") or 0) instruction = normalize_inline_whitespace(str(item.get("instruction") or "")) lines.append(f"{field_name} | {field_type} | {document_count} | {instruction}") return "\n".join(lines) def clear_session_browsing(session_state: dict[str, object]) -> dict[str, object]: session_state["browsing"] = coerce_mode_payloads({}, coerce_browsing_payload) return session_state def persist_session_state(paths: dict[str, Path], session_state: dict[str, object]) -> dict[str, object]: write_session_state(paths, session_state) return read_session_state(paths) def persist_scope_to_session( paths: dict[str, Path], scope: dict[str, object], *, reset_browsing: bool = True, ) -> dict[str, object]: session_state = read_session_state(paths) normalized_scope = coerce_scope_payload(scope) if normalized_scope: normalized_scope["set_at"] = utc_now() session_state["scope"] = normalized_scope if reset_browsing: clear_session_browsing(session_state) return persist_session_state(paths, session_state) def find_saved_scope_name(saved_scopes_state: dict[str, object], requested_name: str) -> str | None: scopes = saved_scopes_state.get("scopes") if not isinstance(scopes, dict): return None normalized_name = normalize_saved_scope_name(requested_name) for existing_name in scopes: if normalize_saved_scope_name(str(existing_name)) == normalized_name: return str(existing_name) return None def save_named_scope( paths: dict[str, Path], scope_name: str, scope: dict[str, object], ) -> dict[str, object]: normalized_scope_name = normalize_inline_whitespace(scope_name) if not normalized_scope_name: raise RetrieverError("Saved scope name cannot be empty.") saved_scopes_state = read_saved_scopes_state(paths) scopes = saved_scopes_state.setdefault("scopes", {}) assert isinstance(scopes, dict) existing_name = find_saved_scope_name(saved_scopes_state, normalized_scope_name) if existing_name is None and len(scopes) >= MAX_SAVED_SCOPES: raise RetrieverError(f"Saved scopes are capped at {MAX_SAVED_SCOPES} per workspace.") if existing_name is not None and existing_name != normalized_scope_name: scopes.pop(existing_name, None) payload = coerce_scope_payload(scope) payload.pop("set_at", None) if payload: payload["saved_at"] = utc_now() scopes[normalized_scope_name] = payload write_saved_scopes_state(paths, saved_scopes_state) return {"status": "ok", "name": normalized_scope_name, "scope": payload} def parse_bates_scope_input(raw_value: str) -> dict[str, str]: begin, end = parse_bates_query(raw_value) if begin is None or end is None: raise RetrieverError("Expected a Bates token or Bates range.") begin_parsed = parse_bates_identifier(begin) end_parsed = parse_bates_identifier(end) if not bates_range_compatible(begin_parsed, end_parsed): raise RetrieverError("Mixed-prefix Bates ranges are not supported; use two separate queries.") return {"begin": begin, "end": end} def intersect_bates_scopes(current_bates: object, incoming_bates: dict[str, str]) -> dict[str, str]: if not isinstance(current_bates, dict): return incoming_bates current_begin = parse_bates_identifier(current_bates.get("begin")) current_end = parse_bates_identifier(current_bates.get("end")) next_begin = parse_bates_identifier(incoming_bates.get("begin")) next_end = parse_bates_identifier(incoming_bates.get("end")) if not all((current_begin, current_end, next_begin, next_end)): return incoming_bates if not bates_range_compatible(current_begin, current_end) or not bates_range_compatible(next_begin, next_end): return incoming_bates if not bates_range_compatible(current_begin, next_begin): raise RetrieverError("Cannot intersect Bates ranges from different series.") overlap_begin = max(int(current_begin["number"]), int(next_begin["number"])) overlap_end = min(int(current_end["number"]), int(next_end["number"])) if overlap_begin > overlap_end: raise RetrieverError("The requested Bates range does not overlap the current Bates scope.") prefix = str(current_begin["prefix"]) width = int(current_begin["width"]) return { "begin": f"{prefix}{overlap_begin:0{width}d}", "end": f"{prefix}{overlap_end:0{width}d}", } def refresh_dataset_name_refs_in_scope(scope: dict[str, object], dataset_id: int, dataset_name: str) -> dict[str, object]: dataset_entries = coerce_scope_dataset_entries(scope.get("dataset")) if not dataset_entries: return scope updated_entries: list[dict[str, object]] = [] for entry in dataset_entries: if int(entry["id"]) == dataset_id: updated_entries.append({"id": dataset_id, "name": dataset_name}) else: updated_entries.append(entry) scope["dataset"] = updated_entries return scope def refresh_dataset_name_refs_in_saved_scopes(paths: dict[str, Path], dataset_id: int, dataset_name: str) -> None: saved_scopes_state = read_saved_scopes_state(paths) scopes = saved_scopes_state.get("scopes") if not isinstance(scopes, dict): return changed = False for scope_payload in scopes.values(): if not isinstance(scope_payload, dict): continue before = json.dumps(scope_payload, ensure_ascii=True, sort_keys=True) refresh_dataset_name_refs_in_scope(scope_payload, dataset_id, dataset_name) after = json.dumps(scope_payload, ensure_ascii=True, sort_keys=True) if before != after: changed = True if changed: write_saved_scopes_state(paths, saved_scopes_state) def filter_expression_references_field(expression: str, field_name: str) -> bool: tokens = tokenize_sql_filter_expression(expression) for token in tokens: if token["kind"] != "identifier": continue if str(token["value"]) == field_name: return True return False def rewrite_filter_expression_field_name( expression: str, old_field_name: str, new_field_name: str, ) -> tuple[str, bool]: tokens = tokenize_sql_filter_expression(expression) parts: list[str] = [] cursor = 0 changed = False for token in tokens: if token["kind"] == "eof": break start = int(token["start"]) end = int(token["end"]) parts.append(expression[cursor:start]) if token["kind"] == "identifier" and str(token["value"]) == old_field_name: parts.append(new_field_name) changed = True else: parts.append(expression[start:end]) cursor = end parts.append(expression[cursor:]) return "".join(parts), changed def field_filter_blocker(scope_name: str, field_name: str, expression: str, *, invalid: bool = False) -> dict[str, object]: if invalid: return { "kind": "invalid_filter_reference", "scope": scope_name, "field_name": field_name, "expression": expression, "message": f"{scope_name} filter could not be rewritten safely for field '{field_name}'.", } return { "kind": "filter_reference", "scope": scope_name, "field_name": field_name, "expression": expression, "message": f"{scope_name} filter references field '{field_name}'.", } def rewrite_scope_filter_field_name( scope_payload: object, old_field_name: str, new_field_name: str, *, scope_name: str, ) -> tuple[dict[str, object], list[dict[str, object]], int]: scope = coerce_scope_payload(scope_payload) expression = normalize_inline_whitespace(str(scope.get("filter") or "")) if not expression: return scope, [], 0 try: next_expression, changed = rewrite_filter_expression_field_name(expression, old_field_name, new_field_name) except RetrieverError: if old_field_name in expression: return scope, [field_filter_blocker(scope_name, old_field_name, expression, invalid=True)], 0 return scope, [], 0 if not changed: return scope, [], 0 scope["filter"] = next_expression return scope, [], 1 def detect_scope_filter_field_refs( scope_payload: object, field_name: str, *, scope_name: str, ) -> list[dict[str, object]]: scope = coerce_scope_payload(scope_payload) expression = normalize_inline_whitespace(str(scope.get("filter") or "")) if not expression: return [] try: if filter_expression_references_field(expression, field_name): return [field_filter_blocker(scope_name, field_name, expression)] except RetrieverError: if field_name in expression: return [field_filter_blocker(scope_name, field_name, expression, invalid=True)] return [] def update_session_display_field_refs( session_state: dict[str, object], old_field_name: str, *, new_field_name: str | None = None, ) -> int: display_root = session_state.get("display") if not isinstance(display_root, dict): display_root = {} changed = 0 for browse_mode in (BROWSE_MODE_DOCUMENTS, BROWSE_MODE_CONVERSATIONS): display_state = session_display_state(session_state, browse_mode=browse_mode) columns = display_state.get("columns") if not isinstance(columns, list): continue next_columns: list[str] = [] seen_names: set[str] = set() branch_changed = False for raw_column in columns: column_name = normalize_inline_whitespace(str(raw_column or "")) if not column_name: branch_changed = True continue if column_name == old_field_name: branch_changed = True if new_field_name is None: continue column_name = new_field_name if column_name in seen_names: branch_changed = True continue seen_names.add(column_name) next_columns.append(column_name) if not branch_changed: continue changed += 1 if next_columns: display_state["columns"] = next_columns else: display_state.pop("columns", None) display_root[browse_mode] = coerce_display_payload(display_state) session_state["display"] = coerce_mode_payloads(display_root, coerce_display_payload) return changed def update_session_sort_field_refs( session_state: dict[str, object], old_field_name: str, *, new_field_name: str | None = None, ) -> int: browsing_root = session_state.get("browsing") if not isinstance(browsing_root, dict): browsing_root = {} changed = 0 for browse_mode in (BROWSE_MODE_DOCUMENTS, BROWSE_MODE_CONVERSATIONS): browsing_state = session_browsing_state(session_state, browse_mode=browse_mode) sort_specs = coerce_sort_specs(browsing_state.get("sort")) if not sort_specs: continue next_specs: list[tuple[str, str]] = [] branch_changed = False seen_specs: set[tuple[str, str]] = set() for field_name, direction in sort_specs: next_field_name = field_name if field_name == old_field_name: branch_changed = True if new_field_name is None: continue next_field_name = new_field_name spec = (next_field_name, direction) if spec in seen_specs: branch_changed = True continue seen_specs.add(spec) next_specs.append(spec) if not branch_changed: continue changed += 1 if next_specs: browsing_state["sort"] = serialize_sort_specs(next_specs) else: browsing_state.pop("sort", None) browsing_root[browse_mode] = coerce_browsing_payload(browsing_state) session_state["browsing"] = coerce_mode_payloads(browsing_root, coerce_browsing_payload) return changed def plan_field_rename_state_changes( paths: dict[str, Path], old_field_name: str, new_field_name: str, ) -> dict[str, object]: session_state = read_session_state(paths) saved_scopes_state = read_saved_scopes_state(paths) blockers: list[dict[str, object]] = [] changes = { "display_columns_updated": update_session_display_field_refs( session_state, old_field_name, new_field_name=new_field_name, ), "sort_specs_updated": update_session_sort_field_refs( session_state, old_field_name, new_field_name=new_field_name, ), "active_scope_filters_updated": 0, "saved_scope_filters_updated": 0, } next_scope, scope_blockers, active_scope_updates = rewrite_scope_filter_field_name( session_state.get("scope"), old_field_name, new_field_name, scope_name="active scope", ) blockers.extend(scope_blockers) changes["active_scope_filters_updated"] = active_scope_updates session_state["scope"] = next_scope scopes = saved_scopes_state.get("scopes") if isinstance(scopes, dict): updated_saved_scope_filters = 0 for scope_name, scope_payload in scopes.items(): if not isinstance(scope_payload, dict): continue next_saved_scope, saved_scope_blockers, saved_scope_updates = rewrite_scope_filter_field_name( scope_payload, old_field_name, new_field_name, scope_name=f"saved scope '{scope_name}'", ) blockers.extend(saved_scope_blockers) scopes[scope_name] = coerce_saved_scope_payload(next_saved_scope) updated_saved_scope_filters += saved_scope_updates changes["saved_scope_filters_updated"] = updated_saved_scope_filters return { "session_state": session_state, "saved_scopes_state": saved_scopes_state, "blockers": blockers, "changes": changes, } def plan_field_delete_state_changes(paths: dict[str, Path], field_name: str) -> dict[str, object]: session_state = read_session_state(paths) saved_scopes_state = read_saved_scopes_state(paths) blockers: list[dict[str, object]] = [] changes = { "display_columns_updated": update_session_display_field_refs(session_state, field_name), "sort_specs_updated": update_session_sort_field_refs(session_state, field_name), "active_scope_filters_blocked": 0, "saved_scope_filters_blocked": 0, } active_scope_blockers = detect_scope_filter_field_refs( session_state.get("scope"), field_name, scope_name="active scope", ) blockers.extend(active_scope_blockers) changes["active_scope_filters_blocked"] = len(active_scope_blockers) scopes = saved_scopes_state.get("scopes") if isinstance(scopes, dict): saved_scope_blockers: list[dict[str, object]] = [] for scope_name, scope_payload in scopes.items(): if not isinstance(scope_payload, dict): continue saved_scope_blockers.extend( detect_scope_filter_field_refs( scope_payload, field_name, scope_name=f"saved scope '{scope_name}'", ) ) blockers.extend(saved_scope_blockers) changes["saved_scope_filters_blocked"] = len(saved_scope_blockers) return { "session_state": session_state, "saved_scopes_state": saved_scopes_state, "blockers": blockers, "changes": changes, } def apply_field_state_change_plan(paths: dict[str, Path], plan: dict[str, object]) -> None: session_state = plan.get("session_state") if isinstance(session_state, dict): write_session_state(paths, session_state) saved_scopes_state = plan.get("saved_scopes_state") if isinstance(saved_scopes_state, dict): write_saved_scopes_state(paths, saved_scopes_state) def split_quoted_comma_values(raw_text: str) -> list[str]: values: list[str] = [] current: list[str] = [] quote_char: str | None = None escaped = False for char in raw_text: if escaped: current.append(char) escaped = False continue if char == "\\": escaped = True continue if quote_char is not None: if char == quote_char: quote_char = None else: current.append(char) continue if char in {"'", '"'}: quote_char = char continue if char == ",": value = normalize_inline_whitespace("".join(current)) if value: values.append(value) current = [] continue current.append(char) if quote_char is not None: raise RetrieverError("Unterminated quote in slash command.") if escaped: current.append("\\") value = normalize_inline_whitespace("".join(current)) if value: values.append(value) return values def shlex_split_slash_tail(raw_tail: str) -> list[str]: try: return shlex.split(raw_tail, posix=True) except ValueError as exc: raise RetrieverError(f"Could not parse slash command arguments: {exc}") from exc def resolve_scope_dataset_selection(connection: sqlite3.Connection, raw_values: list[str]) -> list[dict[str, object]]: if len(raw_values) > MAX_SCOPE_DATASETS: raise RetrieverError(f"Scope datasets are capped at {MAX_SCOPE_DATASETS} entries.") resolved_entries: list[dict[str, object]] = [] seen_ids: set[int] = set() for raw_value in raw_values: matches = find_dataset_rows_by_name(connection, raw_value) if not matches: suggestions = scope_dataset_name_suggestions(connection, raw_value) suggestion_text = f" Did you mean: {', '.join(suggestions)}?" if suggestions else "" raise RetrieverError(f"Unknown dataset name: {raw_value}.{suggestion_text}") row = matches[0] dataset_id = int(row["id"]) if dataset_id in seen_ids: continue seen_ids.add(dataset_id) resolved_entries.append({"id": dataset_id, "name": str(row["dataset_name"])}) return resolved_entries def scope_selector_instances(raw_selector: object) -> list[dict[str, object]]: if isinstance(raw_selector, dict) and isinstance(raw_selector.get("all_of"), list): instances: list[dict[str, object]] = [] for item in raw_selector["all_of"]: normalized = coerce_scope_payload(item) if normalized: instances.append(normalized) return instances normalized = coerce_scope_payload(raw_selector) return [normalized] if normalized else [] def compose_scope_selectors_and(*raw_selectors: object) -> dict[str, object]: instances: list[dict[str, object]] = [] for raw_selector in raw_selectors: instances.extend(scope_selector_instances(raw_selector)) if not instances: return {} if len(instances) == 1: return instances[0] return {"all_of": instances} def preferred_scope_selector_from_run_id(raw_selector: object) -> int | None: preferred_run_id: int | None = None for scope in scope_selector_instances(raw_selector): if scope.get("from_run_id") is not None: preferred_run_id = int(scope["from_run_id"]) return preferred_run_id def and_compose_scope_text(existing_value: object, incoming_value: object) -> str: existing_text = normalize_inline_whitespace(str(existing_value or "")) incoming_text = normalize_inline_whitespace(str(incoming_value or "")) if not existing_text: return incoming_text if not incoming_text: return existing_text return f"({existing_text}) AND ({incoming_text})" def merge_scope_with_search_inputs( raw_scope: object, query: str, raw_filters: list[list[str]] | None, ) -> dict[str, object]: merged_scope = coerce_scope_payload(raw_scope) incoming_scope = derive_search_scope(query, raw_filters) incoming_keyword = normalize_inline_whitespace(str(incoming_scope.get("keyword") or "")) if incoming_keyword: merged_scope["keyword"] = and_compose_scope_text(merged_scope.get("keyword"), incoming_keyword) incoming_bates = incoming_scope.get("bates") if isinstance(incoming_bates, dict): merged_scope["bates"] = intersect_bates_scopes(merged_scope.get("bates"), incoming_bates) incoming_filter = normalize_inline_whitespace(str(incoming_scope.get("filter") or "")) if incoming_filter: merged_scope["filter"] = and_compose_scope_text(merged_scope.get("filter"), incoming_filter) return merged_scope def build_explicit_scope_selector( connection: sqlite3.Connection, *, query: str, raw_bates: str | None, raw_filters: list[list[str]] | None, dataset_names: list[str] | None, from_run_id: int | None, ) -> dict[str, object]: selector: dict[str, object] = {} normalized_query = query.strip() if normalized_query: selector["keyword"] = query normalized_bates = normalize_inline_whitespace(str(raw_bates or "")) if normalized_bates: selector["bates"] = parse_bates_scope_input(normalized_bates) expressions = normalize_sql_filter_expressions(raw_filters) if expressions: selector["filter"] = " AND ".join(f"({expression})" for expression in expressions) raw_dataset_names = [ normalize_inline_whitespace(str(dataset_name or "")) for dataset_name in (dataset_names or []) ] normalized_dataset_names = [dataset_name for dataset_name in raw_dataset_names if dataset_name] if normalized_dataset_names: selector["dataset"] = resolve_scope_dataset_selection(connection, normalized_dataset_names) resolved_from_run_id = resolve_scope_from_run_id(connection, from_run_id) if resolved_from_run_id is not None: selector["from_run_id"] = resolved_from_run_id return coerce_scope_payload(selector) def build_effective_scope_selector( connection: sqlite3.Connection, paths: dict[str, Path], *, query: str, raw_bates: str | None, raw_filters: list[list[str]] | None, dataset_names: list[str] | None, from_run_id: int | None, select_from_scope: bool, ) -> dict[str, object]: base_scope = {} if select_from_scope: base_scope = coerce_scope_payload(read_session_state(paths).get("scope")) base_scope.pop("set_at", None) explicit_selector = build_explicit_scope_selector( connection, query=query, raw_bates=raw_bates, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, ) return compose_scope_selectors_and(base_scope, explicit_selector) def parse_fill_slash_arguments(normalized_tail: str) -> dict[str, object]: tokens = shlex_split_slash_tail(normalized_tail) if len(tokens) < 2: raise RetrieverError("Usage: /fill [on ] [--confirm]") confirm = False filtered_tokens: list[str] = [] for token in tokens: if token == "--confirm": confirm = True continue filtered_tokens.append(token) tokens = filtered_tokens if len(tokens) < 2: raise RetrieverError("Usage: /fill [on ] [--confirm]") field_name = tokens[0] on_index = -1 for index in range(1, len(tokens)): if tokens[index] == "on": on_index = index break if on_index == -1: value_tokens = tokens[1:] doc_refs: list[str] = [] else: value_tokens = tokens[1:on_index] trailing_tokens = tokens[on_index + 1 :] if not trailing_tokens: raise RetrieverError("Usage: /fill on [--confirm]") raw_doc_ref_text = " ".join(trailing_tokens) doc_refs = split_quoted_comma_values(raw_doc_ref_text) if len(doc_refs) == 1 and "," not in raw_doc_ref_text and len(trailing_tokens) > 1: doc_refs = [ normalize_inline_whitespace(token) for token in trailing_tokens if normalize_inline_whitespace(token) ] if not value_tokens: raise RetrieverError("Usage: /fill [on ] [--confirm]") clear = len(value_tokens) == 1 and value_tokens[0].lower() == "clear" value = None if clear else normalize_inline_whitespace(" ".join(value_tokens)) if not clear and not value: raise RetrieverError("Fill value cannot be empty.") return { "field_name": field_name, "value": value, "clear": clear, "doc_refs": doc_refs, "confirm": confirm, } def resolve_scope_document_search_with_explicit_sort( connection: sqlite3.Connection, raw_scope: object, sort_field: str | None, order: str | None, ) -> dict[str, object]: normalized_sort_field = sort_field normalized_order = (order or "desc").lower() keyword_query = normalize_inline_whitespace(str(coerce_scope_payload(raw_scope).get("keyword") or "")) if sort_field == "relevance" and not keyword_query: raise RetrieverError("Sort 'relevance' requires a non-empty query.") if sort_field and sort_field != "relevance": normalized_sort_field = resolve_sort_field_name(connection, sort_field) selection = resolve_scope_document_search( connection, raw_scope, sort_specs=[(str(normalized_sort_field), normalized_order)], ) selection["sort"] = str(normalized_sort_field) selection["order"] = normalized_order selection["sort_spec"] = f"{normalized_sort_field} {normalized_order}" return selection selection = resolve_scope_document_search(connection, raw_scope) if sort_field == "relevance": selection["sort"] = "relevance" selection["order"] = (order or "asc").lower() selection["sort_spec"] = f"relevance {selection['order']}" return selection def build_entity_header_payload(payload: dict[str, object]) -> dict[str, str]: total_hits = int(payload.get("total_hits") or 0) page = int(payload.get("page") or 1) per_page = int(payload.get("per_page") or DEFAULT_PAGE_SIZE) start_index = 0 if total_hits == 0 else ((page - 1) * per_page) + 1 end_index = 0 if total_hits == 0 else min(total_hits, page * per_page) sort_summary = str(payload.get("sort_spec") or f"{payload.get('sort')} {payload.get('order')}") header: dict[str, str] = {} query = normalize_inline_whitespace(str(payload.get("query") or "")) if query: header["keyword"] = f"Entity query: {query!r}" if bool(payload.get("include_ignored")): header["filters"] = "Include ignored entities" header["sort"] = f"Sort: {sort_summary}" header["page"] = f"Page: {page} of {payload.get('total_pages')} (entities {start_index}-{end_index} of {total_hits})" return header def decorate_entity_list_payload( payload: dict[str, object], column_defs: list[dict[str, str]], *, warnings: list[str] | None = None, ) -> dict[str, object]: per_page = int(payload.get("limit") or DEFAULT_PAGE_SIZE) offset = int(payload.get("offset") or 0) total_hits = int(payload.get("total_hits") or payload.get("total") or 0) total_pages = max(1, (total_hits + per_page - 1) // per_page) page = (offset // per_page) + 1 results = [dict(item) for item in payload.get("entities", []) if isinstance(item, dict)] for item in results: item["entity_status"] = item.get("canonical_status") item["display_values"] = build_summary_display_values(item, column_defs) decorated = { **payload, "browse_mode": BROWSE_MODE_ENTITIES, "page": page, "per_page": per_page, "offset": offset, "total_hits": total_hits, "total_pages": total_pages, "results": results, "entities": results, "filters": [], "scope": {}, "display": build_display_payload(column_defs, per_page), } decorated["header"] = build_entity_header_payload(decorated) if warnings: decorated["warnings"] = warnings decorated["rendered_markdown"] = render_search_markdown(decorated, column_defs) return decorated def entity_browsing_options(session_state: dict[str, object]) -> tuple[str | None, bool]: browsing = session_browsing_state(session_state, browse_mode=BROWSE_MODE_ENTITIES) query = browsing.get("query") normalized_query = normalize_whitespace(str(query or "")) if isinstance(query, str) else "" return (normalized_query or None), bool(browsing.get("include_ignored")) def persist_entity_browsing_result( paths: dict[str, Path], session_state: dict[str, object], payload: dict[str, object], sort_specs: list[tuple[str, str]] | None, ) -> dict[str, object]: session_state["browse_mode"] = BROWSE_MODE_ENTITIES browsing_root = session_state.get("browsing") if not isinstance(browsing_root, dict): browsing_root = {} browsing_payload: dict[str, object] = { "offset": int(payload.get("offset") or 0), "total_known": int(payload.get("total_hits") or 0), "include_ignored": bool(payload.get("include_ignored")), "run_at": utc_now(), } query = normalize_whitespace(str(payload.get("query") or "")) if query: browsing_payload["query"] = query if sort_specs: browsing_payload["sort"] = serialize_sort_specs(sort_specs) browsing_root[BROWSE_MODE_ENTITIES] = coerce_browsing_payload(browsing_payload) session_state["browsing"] = coerce_mode_payloads(browsing_root, coerce_browsing_payload) persist_session_state(paths, session_state) return payload def run_entity_browsing_from_session( root: Path, paths: dict[str, Path], session_state: dict[str, object] | None = None, *, offset: int | None = None, sort_specs: list[tuple[str, str]] | None = None, ) -> dict[str, object]: normalized_session_state = read_session_state(paths) if session_state is None else session_state normalized_session_state["browse_mode"] = BROWSE_MODE_ENTITIES effective_sort_specs = ( sort_specs if sort_specs is not None else session_sort_specs(normalized_session_state, browse_mode=BROWSE_MODE_ENTITIES) ) current_offset = int( session_browsing_state(normalized_session_state, browse_mode=BROWSE_MODE_ENTITIES).get("offset") or 0 ) query, include_ignored = entity_browsing_options(normalized_session_state) connection = connect_db(paths["db_path"]) try: display_column_defs, display_warnings, normalized_session_state = resolve_session_display_columns( connection, paths, normalized_session_state, browse_mode=BROWSE_MODE_ENTITIES, ) finally: connection.close() per_page = session_page_size(normalized_session_state, browse_mode=BROWSE_MODE_ENTITIES) requested_offset = current_offset if offset is None else offset payload = list_entities( root, query=query, limit=per_page, offset=requested_offset, sort_specs=effective_sort_specs or None, include_ignored=include_ignored, ) total_hits = int(payload.get("total_hits") or 0) total_pages = max(1, (total_hits + per_page - 1) // per_page) if total_hits > 0 and requested_offset >= total_hits: requested_offset = (total_pages - 1) * per_page payload = list_entities( root, query=query, limit=per_page, offset=requested_offset, sort_specs=effective_sort_specs or None, include_ignored=include_ignored, ) decorated = decorate_entity_list_payload(payload, display_column_defs, warnings=display_warnings) return persist_entity_browsing_result( paths, normalized_session_state, decorated, effective_sort_specs or None, ) def run_entity_list_command( root: Path, *, query: str | None, limit: int, offset: int, sort: str | None, order: str | None, include_ignored: bool, ) -> dict[str, object]: sort_specs = normalize_entity_list_sort_specs( sort=sort or ("document_count" if order else None), order=order, ) if (sort or order) else None payload = list_entities( root, query=query, limit=limit, offset=offset, sort_specs=sort_specs, include_ignored=include_ignored, ) paths = workspace_paths(root) ensure_layout(paths) session_state = read_session_state(paths) display_root = session_state.get("display") if not isinstance(display_root, dict): display_root = {} display_state = session_display_state(session_state, browse_mode=BROWSE_MODE_ENTITIES) page_size = min(int(payload.get("limit") or DEFAULT_PAGE_SIZE), MAX_PAGE_SIZE) if page_size == DEFAULT_PAGE_SIZE: display_state.pop("page_size", None) else: display_state["page_size"] = page_size display_root[BROWSE_MODE_ENTITIES] = coerce_display_payload(display_state) session_state["display"] = coerce_mode_payloads(display_root, coerce_display_payload) persist_entity_browsing_result(paths, session_state, payload, sort_specs) return payload def normalize_conversation_list_sort_specs( connection: sqlite3.Connection, *, sort: str | None, order: str | None, ) -> list[tuple[str, str]] | None: if not sort and not order: return None field_name = resolve_sort_field_name( connection, sort or "last_activity", browse_mode=BROWSE_MODE_CONVERSATIONS, ) direction = normalize_inline_whitespace(str(order or "desc")).lower() if direction not in {"asc", "desc"}: raise RetrieverError("Sort direction must be 'asc' or 'desc'.") return [(field_name, direction)] def list_conversations( root: Path, *, query: str | None = None, raw_filters: list[list[str]] | None = None, dataset_names: list[str] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, limit: int = 50, offset: int = 0, sort: str | None = None, order: str | None = None, ) -> dict[str, object]: normalized_limit = max(1, min(int(limit or 50), MAX_PAGE_SIZE)) normalized_offset = max(0, int(offset or 0)) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) selector = build_effective_scope_selector( connection, paths, query=normalize_inline_whitespace(str(query or "")), raw_bates=None, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) sort_specs = normalize_conversation_list_sort_specs(connection, sort=sort, order=order) session_state = read_session_state(paths) display_column_defs, display_warnings, session_state = resolve_session_display_columns( connection, paths, session_state, browse_mode=BROWSE_MODE_CONVERSATIONS, ) finally: connection.close() payload = search_with_scope( root, selector, per_page=normalized_limit, offset=normalized_offset, sort_specs=sort_specs, display_column_defs=display_column_defs, warnings=display_warnings, browse_mode=BROWSE_MODE_CONVERSATIONS, ) payload["limit"] = normalized_limit payload["conversations"] = payload["results"] persisted_session_state = persist_display_preferences( paths, session_state, display_column_defs, normalized_limit, browse_mode=BROWSE_MODE_CONVERSATIONS, ) persist_browsing_search_result( paths, persisted_session_state, payload, sort_specs, browse_mode=BROWSE_MODE_CONVERSATIONS, ) return payload def persist_browsing_search_result( paths: dict[str, Path], session_state: dict[str, object], payload: dict[str, object], sort_specs: list[tuple[str, str]] | None, *, browse_mode: str | None = None, ) -> dict[str, object]: effective_browse_mode = normalize_browse_mode( browse_mode or payload.get("browse_mode") or session_browse_mode(session_state) ) session_state["scope"] = coerce_scope_payload(payload.get("scope")) session_state["browse_mode"] = effective_browse_mode browsing_root = session_state.get("browsing") if not isinstance(browsing_root, dict): browsing_root = {} browsing_payload: dict[str, object] = { "offset": int(payload.get("offset") or 0), "total_known": int(payload.get("total_hits") or 0), "run_at": utc_now(), } if sort_specs: browsing_payload["sort"] = serialize_sort_specs(sort_specs) browsing_root[effective_browse_mode] = coerce_browsing_payload(browsing_payload) session_state["browsing"] = coerce_mode_payloads(browsing_root, coerce_browsing_payload) persist_session_state(paths, session_state) return payload def persist_direct_view_search_result( root: Path, payload: dict[str, object], column_defs: list[dict[str, str]], *, sort_specs: list[tuple[str, str]] | None, browse_mode: str = BROWSE_MODE_DOCUMENTS, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) session_state = read_session_state(paths) persisted_session_state = persist_display_preferences( paths, session_state, column_defs, int(payload.get("per_page") or DEFAULT_PAGE_SIZE), browse_mode=browse_mode, ) persist_browsing_search_result( paths, persisted_session_state, payload, sort_specs, browse_mode=browse_mode, ) return payload def run_browsing_search_from_session( root: Path, paths: dict[str, Path], session_state: dict[str, object] | None = None, *, offset: int | None = None, sort_specs: list[tuple[str, str]] | None = None, browse_mode: str | None = None, ) -> dict[str, object]: normalized_session_state = read_session_state(paths) if session_state is None else session_state effective_browse_mode = normalize_browse_mode(browse_mode or session_browse_mode(normalized_session_state)) if effective_browse_mode == BROWSE_MODE_ENTITIES: return run_entity_browsing_from_session( root, paths, normalized_session_state, offset=offset, sort_specs=sort_specs, ) normalized_session_state["browse_mode"] = effective_browse_mode effective_sort_specs = ( sort_specs if sort_specs is not None else session_sort_specs(normalized_session_state, browse_mode=effective_browse_mode) ) current_offset = int( session_browsing_state(normalized_session_state, browse_mode=effective_browse_mode).get("offset") or 0 ) connection = connect_db(paths["db_path"]) try: display_column_defs, display_warnings, normalized_session_state = resolve_session_display_columns( connection, paths, normalized_session_state, browse_mode=effective_browse_mode, ) finally: connection.close() payload = search_with_scope( root, normalized_session_state.get("scope", {}), per_page=session_page_size(normalized_session_state, browse_mode=effective_browse_mode), offset=current_offset if offset is None else offset, sort_specs=effective_sort_specs or None, display_column_defs=display_column_defs, warnings=display_warnings, browse_mode=effective_browse_mode, ) return persist_browsing_search_result( paths, normalized_session_state, payload, effective_sort_specs or None, browse_mode=effective_browse_mode, ) def run_scope_search_from_session(root: Path, paths: dict[str, Path], scope: dict[str, object]) -> dict[str, object]: session_state = persist_scope_to_session(paths, scope) browse_mode = session_browse_mode(session_state) if browse_mode == BROWSE_MODE_ENTITIES: session_state["browse_mode"] = BROWSE_MODE_DOCUMENTS session_state = persist_session_state(paths, session_state) browse_mode = BROWSE_MODE_DOCUMENTS return run_browsing_search_from_session(root, paths, session_state, offset=0, browse_mode=browse_mode) EXPORT_TABLE_ALIASES = { "document": "documents", "documents": "documents", "doc": "documents", "docs": "documents", "entity": "entities", "entities": "entities", "conversation": "conversations", "conversations": "conversations", } SLASH_EXPORT_KNOWN_TABLES = set(EXPORT_TABLE_ALIASES) | {"dataset", "datasets"} SLASH_EXPORT_BOOLEAN_OPTIONS = { "--include-ignored", "--no-scope", "--portable", "--portable-workspace", "--run-to-completion", } SLASH_EXPORT_VALUE_OPTIONS = { "--bates", "--budget-seconds", "--dataset", "--doc-id", "--family-mode", "--field", "--filter", "--from-run-id", "--keyword", "--limit", "--order", "--sort", "--sort-by", "--sort-order", } def normalize_export_table_name(table_name: str) -> str: normalized = normalize_inline_whitespace(str(table_name or "documents")).lower() table = EXPORT_TABLE_ALIASES.get(normalized) if table is None: raise RetrieverError("Supported export tables are: documents, entities, conversations.") return table def slash_export_timestamp() -> str: digits = "".join(character for character in utc_now() if character.isdigit()) return digits[:14] or "export" def slash_export_default_output_path(prefix: str, suffix: str) -> str: return f"{prefix}-{slash_export_timestamp()}.{suffix}" def parse_slash_export_options(tokens: list[str]) -> tuple[list[str], dict[str, object]]: positionals: list[str] = [] options: dict[str, object] = {} repeat_options = {"--dataset", "--doc-id", "--field", "--filter"} index = 0 while index < len(tokens): token = tokens[index] if not token.startswith("--"): positionals.append(token) index += 1 continue if token in SLASH_EXPORT_BOOLEAN_OPTIONS: options[token] = True index += 1 continue if token not in SLASH_EXPORT_VALUE_OPTIONS: raise RetrieverError(f"Unknown /export option: {token}") index += 1 values: list[str] = [] while index < len(tokens) and not tokens[index].startswith("--"): values.append(tokens[index]) index += 1 if token != "--filter": break if not values: raise RetrieverError(f"{token} requires a value.") value: object = values if token == "--filter" else values[0] canonical_token = { "--sort-by": "--sort", "--sort-order": "--order", "--portable": "--portable-workspace", }.get(token, token) if canonical_token in repeat_options: options.setdefault(canonical_token, []) existing_values = options[canonical_token] if not isinstance(existing_values, list): raise RetrieverError(f"Internal /export parse error for {canonical_token}.") existing_values.append(value) else: options[canonical_token] = value return positionals, options def parse_optional_int_option(options: dict[str, object], option_name: str) -> int | None: if option_name not in options: return None try: return int(str(options[option_name])) except ValueError as exc: raise RetrieverError(f"{option_name} must be an integer.") from exc def normalize_slash_export_filters(options: dict[str, object]) -> list[list[str]] | None: raw_filters = options.get("--filter") if raw_filters is None: return None if not isinstance(raw_filters, list): raise RetrieverError("Internal /export parse error for --filter.") normalized_filters: list[list[str]] = [] for raw_filter in raw_filters: if isinstance(raw_filter, list): normalized_filters.append([str(part) for part in raw_filter]) else: normalized_filters.append([str(raw_filter)]) return normalized_filters def run_export_table_slash_command( root: Path, paths: dict[str, Path], connection: sqlite3.Connection, session_state: dict[str, object], tokens: list[str], ) -> dict[str, object]: positionals, options = parse_slash_export_options(tokens) table_name = "documents" if positionals and positionals[0].lower() in SLASH_EXPORT_KNOWN_TABLES: table_name = normalize_export_table_name(positionals.pop(0)) output_path = slash_export_default_output_path(table_name, "csv") if positionals and (positionals[0].lower().endswith(".csv") or "/" in positionals[0] or "\\" in positionals[0]): output_path = positionals.pop(0) query = " ".join(positionals).strip() if "--keyword" in options: query = str(options["--keyword"]) fields = [str(field) for field in options.get("--field", [])] if "--field" in options else None document_ids = None if "--doc-id" in options: raw_document_ids = options["--doc-id"] if not isinstance(raw_document_ids, list): raise RetrieverError("Internal /export parse error for --doc-id.") document_ids = [] for raw_document_id in raw_document_ids: try: document_ids.append(int(str(raw_document_id))) except ValueError as exc: raise RetrieverError("--doc-id must be an integer.") from exc budget_seconds = parse_optional_int_option(options, "--budget-seconds") limit = parse_optional_int_option(options, "--limit") select_from_scope = not bool(options.get("--no-scope")) and not document_ids payload = export_table_csv_start( root, output_path, table_name=table_name, raw_fields=fields, document_ids=document_ids, query=query, raw_filters=normalize_slash_export_filters(options), sort_field=str(options["--sort"]) if "--sort" in options else None, order=str(options["--order"]) if "--order" in options else None, select_from_scope=select_from_scope, include_ignored_entities=bool(options.get("--include-ignored")), limit=limit, budget_seconds=budget_seconds, run_to_completion=bool(options.get("--run-to-completion")), ) payload["slash_command"] = "/export table" payload["table"] = table_name payload["export_label"] = "table" return payload def run_export_archive_slash_command(root: Path, tokens: list[str]) -> dict[str, object]: positionals, options = parse_slash_export_options(tokens) output_path = slash_export_default_output_path("archive", "zip") if positionals and (positionals[0].lower().endswith(".zip") or "/" in positionals[0] or "\\" in positionals[0]): output_path = positionals.pop(0) query = " ".join(positionals).strip() if "--keyword" in options: query = str(options["--keyword"]) raw_datasets = options.get("--dataset") dataset_names = [str(item) for item in raw_datasets] if isinstance(raw_datasets, list) else None budget_seconds = parse_optional_int_option(options, "--budget-seconds") from_run_id = parse_optional_int_option(options, "--from-run-id") seed_limit = parse_optional_int_option(options, "--limit") payload = export_archive_start( root, output_path, dataset_names=dataset_names, query=query, raw_bates=str(options["--bates"]) if "--bates" in options else None, raw_filters=normalize_slash_export_filters(options), from_run_id=from_run_id, select_from_scope=not bool(options.get("--no-scope")), family_mode=str(options.get("--family-mode") or "exact"), seed_limit=seed_limit, portable_workspace=bool(options.get("--portable-workspace") or options.get("--portable")), budget_seconds=budget_seconds, run_to_completion=bool(options.get("--run-to-completion")), ) payload["slash_command"] = "/export archive" payload["export_label"] = "archive" return payload def run_export_status_slash_command(root: Path, tokens: list[str]) -> dict[str, object]: positionals, options = parse_slash_export_options(tokens) budget_seconds = parse_optional_int_option(options, "--budget-seconds") kind: str | None = None run_id: str | None = None if positionals: first_token = positionals.pop(0) first_token_normalized = first_token.lower() if first_token_normalized in {"table", "csv"}: kind = "csv" elif first_token_normalized == "archive": kind = "archive" else: run_id = first_token if positionals: if run_id is not None: raise RetrieverError("Usage: /export status [table|archive] [run-id] [--budget-seconds N]") run_id = positionals.pop(0) if positionals: raise RetrieverError("Usage: /export status [table|archive] [run-id] [--budget-seconds N]") if kind is not None: payload = export_status(root, export_kind=kind, run_id=run_id, budget_seconds=budget_seconds) payload["slash_command"] = "/export status" payload["export_label"] = "table" if kind == "csv" else kind return payload if run_id is not None: for candidate_kind in ("csv", "archive"): candidate_payload = export_status(root, export_kind=candidate_kind, run_id=run_id, budget_seconds=budget_seconds) if candidate_payload.get("run_id"): candidate_payload["slash_command"] = "/export status" candidate_payload["export_label"] = "table" if candidate_kind == "csv" else candidate_kind return candidate_payload return { "ok": True, "status": "none", "run_id": run_id, "exports": {}, "next_recommended_commands": [], } table_payload = export_status(root, export_kind="csv", budget_seconds=budget_seconds) archive_payload = export_status(root, export_kind="archive", budget_seconds=budget_seconds) active_exports = [ {"export_label": label, "run_id": payload.get("run_id"), "status": payload.get("status"), "phase": payload.get("phase")} for label, payload in (("table", table_payload), ("archive", archive_payload)) if payload.get("status") in EXPORT_RUN_ACTIVE_STATUSES ] next_commands: list[str] = [] for payload in (table_payload, archive_payload): for command in payload.get("next_recommended_commands") or []: if isinstance(command, str) and command not in next_commands: next_commands.append(command) return { "ok": True, "status": "ok", "slash_command": "/export status", "exports": { "table": table_payload, "archive": archive_payload, }, "active_exports": active_exports, "next_recommended_commands": next_commands, } def run_export_slash_command( root: Path, paths: dict[str, Path], connection: sqlite3.Connection, session_state: dict[str, object], normalized_tail: str, ) -> dict[str, object]: tokens = shlex_split_slash_tail(normalized_tail) if normalized_tail else [] if not tokens: raise RetrieverError("Usage: /export ...") subcommand = tokens.pop(0).lower() if subcommand in {"table", "csv"}: return run_export_table_slash_command(root, paths, connection, session_state, tokens) if subcommand == "archive": return run_export_archive_slash_command(root, tokens) if subcommand == "status": return run_export_status_slash_command(root, tokens) if subcommand in {"preview", "previews"}: raise RetrieverError( "/export previews is deferred until preview exports are resumable. " "Use export-previews directly only for small/debug exports." ) raise RetrieverError("Usage: /export ...") def run_slash_command(root: Path, raw_command: str) -> dict[str, object]: command_name, normalized_tail = parse_slash_command_text(raw_command) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) session_state = read_session_state(paths) scope = coerce_scope_payload(session_state.get("scope")) active_browse_mode = session_browse_mode(session_state) if command_name == "export": return run_export_slash_command(root, paths, connection, session_state, normalized_tail) if command_name == "scope": scope_args = shlex_split_slash_tail(normalized_tail) if normalized_tail else [] if not scope_args: return {"status": "ok", "scope": scope} subcommand = scope_args[0] if subcommand == "list": if len(scope_args) != 1: raise RetrieverError("Usage: /scope list") return {"status": "ok", "saved_scopes": saved_scope_summaries(paths)} if subcommand == "clear": return run_scope_search_from_session(root, paths, {}) if subcommand == "save": if len(scope_args) != 2: raise RetrieverError("Usage: /scope save ") return save_named_scope(paths, scope_args[1], scope) if subcommand == "load": if len(scope_args) != 2: raise RetrieverError("Usage: /scope load ") saved_scopes_state = read_saved_scopes_state(paths) existing_name = find_saved_scope_name(saved_scopes_state, scope_args[1]) if existing_name is None: raise RetrieverError(f"Unknown saved scope: {scope_args[1]}") saved_scope = saved_scopes_state["scopes"][existing_name] loaded_scope = coerce_scope_payload(saved_scope) return run_scope_search_from_session(root, paths, loaded_scope) raise RetrieverError(f"Unknown /scope command: {subcommand}") if command_name == "field": field_args = shlex_split_slash_tail(normalized_tail) if normalized_tail else [] if not field_args or field_args == ["list"]: return list_fields(root) subcommand = field_args[0] if subcommand == "add": if len(field_args) < 3: raise RetrieverError("Usage: /field add [description]") instruction = " ".join(field_args[3:]) if len(field_args) > 3 else None return add_field(root, field_args[1], field_args[2], instruction) if subcommand == "rename": if len(field_args) != 3: raise RetrieverError("Usage: /field rename ") return rename_field(root, field_args[1], field_args[2]) if subcommand == "delete": if len(field_args) < 2: raise RetrieverError("Usage: /field delete [--confirm]") confirm = False extra_tokens: list[str] = [] for token in field_args[2:]: if token == "--confirm": confirm = True else: extra_tokens.append(token) if extra_tokens: raise RetrieverError("Usage: /field delete [--confirm]") return delete_field(root, field_args[1], confirm=confirm) if subcommand == "describe": if len(field_args) < 2: raise RetrieverError("Usage: /field describe | /field describe --clear") clear = False text_tokens: list[str] = [] for token in field_args[2:]: if token == "--clear": clear = True else: text_tokens.append(token) if clear and text_tokens: raise RetrieverError("Usage: /field describe | /field describe --clear") if not clear and not text_tokens: raise RetrieverError("Usage: /field describe | /field describe --clear") return describe_field( root, field_args[1], text=None if clear else " ".join(text_tokens), clear=clear, ) if subcommand == "type": if len(field_args) != 3: raise RetrieverError("Usage: /field type ") return change_field_type(root, field_args[1], field_args[2]) raise RetrieverError(f"Unknown /field command: {subcommand}") if command_name == "fill": if not normalized_tail: raise RetrieverError("Usage: /fill [on ] [--confirm]") fill_args = parse_fill_slash_arguments(normalized_tail) raw_doc_refs = fill_args["doc_refs"] if isinstance(raw_doc_refs, list) and raw_doc_refs: document_ids = resolve_fill_document_refs(connection, [str(item) for item in raw_doc_refs]) return fill_field( root, field_name=str(fill_args["field_name"]), value=(str(fill_args["value"]) if fill_args["value"] is not None else None), clear=bool(fill_args["clear"]), document_ids=document_ids, confirm=bool(fill_args["confirm"]), ) return fill_field( root, field_name=str(fill_args["field_name"]), value=(str(fill_args["value"]) if fill_args["value"] is not None else None), clear=bool(fill_args["clear"]), select_from_scope=True, confirm=bool(fill_args["confirm"]), ) if command_name in {"documents", "conversations", "entities"}: if normalized_tail: raise RetrieverError(f"Usage: /{command_name}") if command_name == "conversations": target_browse_mode = BROWSE_MODE_CONVERSATIONS elif command_name == "entities": target_browse_mode = BROWSE_MODE_ENTITIES else: target_browse_mode = BROWSE_MODE_DOCUMENTS updated_session_state = read_session_state(paths) updated_session_state["browse_mode"] = target_browse_mode updated_session_state = persist_session_state(paths, updated_session_state) return run_browsing_search_from_session( root, paths, updated_session_state, browse_mode=target_browse_mode, ) if command_name == "page-size": if not normalized_tail: return { "status": "ok", "browse_mode": active_browse_mode, "page_size": session_page_size(session_state, browse_mode=active_browse_mode), } updated_session_state = read_session_state(paths) effective_browse_mode = session_browse_mode(updated_session_state) display_root = updated_session_state.get("display") if not isinstance(display_root, dict): display_root = {} display_state = session_display_state(updated_session_state, browse_mode=effective_browse_mode) display_state["page_size"] = parse_page_size_value(normalized_tail) display_root[effective_browse_mode] = coerce_display_payload(display_state) updated_session_state["display"] = coerce_mode_payloads(display_root, coerce_display_payload) updated_session_state = persist_session_state(paths, updated_session_state) return run_browsing_search_from_session( root, paths, updated_session_state, browse_mode=effective_browse_mode, ) if command_name == "columns": updated_session_state = read_session_state(paths) effective_browse_mode = session_browse_mode(updated_session_state) if not normalized_tail: column_defs, warnings, updated_session_state = resolve_session_display_columns( connection, paths, updated_session_state, browse_mode=effective_browse_mode, ) payload: dict[str, object] = { "status": "ok", "browse_mode": effective_browse_mode, "display": build_display_payload( column_defs, session_page_size(updated_session_state, browse_mode=effective_browse_mode), ), } if warnings: payload["warnings"] = warnings return payload if normalized_tail == "list": return { "status": "ok", "browse_mode": effective_browse_mode, "columns": displayable_field_entries(connection, browse_mode=effective_browse_mode), } if normalized_tail == "default": display_root = updated_session_state.get("display") if not isinstance(display_root, dict): display_root = {} display_state = session_display_state(updated_session_state, browse_mode=effective_browse_mode) display_state.pop("columns", None) display_root[effective_browse_mode] = coerce_display_payload(display_state) updated_session_state["display"] = coerce_mode_payloads(display_root, coerce_display_payload) updated_session_state = persist_session_state(paths, updated_session_state) return run_browsing_search_from_session( root, paths, updated_session_state, browse_mode=effective_browse_mode, ) subcommand, _, remainder = normalized_tail.partition(" ") if subcommand not in {"set", "add", "remove"}: raise RetrieverError("Usage: /columns, /columns set , /columns add , /columns remove , or /columns default") if not remainder.strip(): raise RetrieverError("Column selection cannot be empty.") current_column_defs, _, updated_session_state = resolve_session_display_columns( connection, paths, updated_session_state, browse_mode=effective_browse_mode, ) current_columns = display_column_names(current_column_defs) if subcommand == "set": requested_columns = parse_display_columns_argument(remainder.strip()) next_column_defs, _, _ = resolve_display_column_definitions( connection, requested_columns, drop_missing=False, browse_mode=effective_browse_mode, ) elif subcommand == "add": field_name = normalize_inline_whitespace(remainder.strip()) next_column_defs, _, _ = resolve_display_column_definitions( connection, current_columns + [field_name], drop_missing=False, browse_mode=effective_browse_mode, ) else: field_name = normalize_inline_whitespace(remainder.strip()) field_def = resolve_browse_field_definition(connection, field_name, browse_mode=effective_browse_mode) canonical_name = str(field_def["field_name"]) if canonical_name not in current_columns: raise RetrieverError(f"Column '{canonical_name}' is not in the current display set.") remaining_columns = [column_name for column_name in current_columns if column_name != canonical_name] if not remaining_columns: raise RetrieverError("Display must include at least one column. Use /columns default to reset.") next_column_defs, _, _ = resolve_display_column_definitions( connection, remaining_columns, drop_missing=False, browse_mode=effective_browse_mode, ) updated_session_state = persist_display_columns( paths, updated_session_state, next_column_defs, browse_mode=effective_browse_mode, ) return run_browsing_search_from_session( root, paths, updated_session_state, browse_mode=effective_browse_mode, ) if command_name == "sort": updated_session_state = read_session_state(paths) effective_browse_mode = session_browse_mode(updated_session_state) if not normalized_tail: return { "status": "ok", "browse_mode": effective_browse_mode, **active_sort_payload(scope, updated_session_state, browse_mode=effective_browse_mode), } if normalized_tail == "list": return { "status": "ok", "browse_mode": effective_browse_mode, "sortable_fields": sortable_field_entries(connection, browse_mode=effective_browse_mode), } if normalized_tail == "default": return run_browsing_search_from_session( root, paths, updated_session_state, offset=0, sort_specs=[], browse_mode=effective_browse_mode, ) sort_specs = parse_slash_sort_specs(connection, normalized_tail, browse_mode=effective_browse_mode) return run_browsing_search_from_session( root, paths, updated_session_state, offset=0, sort_specs=sort_specs, browse_mode=effective_browse_mode, ) if command_name in {"next", "previous"}: updated_session_state = read_session_state(paths) effective_browse_mode = session_browse_mode(updated_session_state) per_page = session_page_size(updated_session_state, browse_mode=effective_browse_mode) current_offset = int( session_browsing_state(updated_session_state, browse_mode=effective_browse_mode).get("offset") or 0 ) target_offset = current_offset + per_page if command_name == "next" else max(0, current_offset - per_page) return run_browsing_search_from_session( root, paths, updated_session_state, offset=target_offset, browse_mode=effective_browse_mode, ) if command_name == "page": updated_session_state = read_session_state(paths) effective_browse_mode = session_browse_mode(updated_session_state) if not normalized_tail: return { "status": "ok", **active_page_payload(updated_session_state, browse_mode=effective_browse_mode), } per_page = session_page_size(updated_session_state, browse_mode=effective_browse_mode) current_offset = int( session_browsing_state(updated_session_state, browse_mode=effective_browse_mode).get("offset") or 0 ) page_token = normalize_inline_whitespace(normalized_tail).lower() if page_token == "first": target_offset = 0 elif page_token == "last": target_offset = 10**12 elif page_token == "next": target_offset = current_offset + per_page elif page_token == "previous": target_offset = max(0, current_offset - per_page) else: try: page_number = int(page_token) except ValueError as exc: raise RetrieverError("Usage: /page ") from exc if page_number < 1: raise RetrieverError("Page number must be >= 1.") target_offset = (page_number - 1) * per_page return run_browsing_search_from_session( root, paths, updated_session_state, offset=target_offset, browse_mode=effective_browse_mode, ) if command_name == "search": if not normalized_tail: return {"status": "ok", "keyword": normalize_inline_whitespace(str(scope.get("keyword") or "")) or None} if normalized_tail == "clear": scope.pop("keyword", None) scope.pop("bates", None) return run_scope_search_from_session(root, paths, scope) force_fts = False within = False query_text = normalized_tail if normalized_tail.startswith("--within "): within = True query_text = normalized_tail[len("--within "):].lstrip() elif normalized_tail.startswith("--fts "): force_fts = True query_text = normalized_tail[len("--fts "):].lstrip() if not query_text: raise RetrieverError("Search text cannot be empty.") incoming_bates = None if force_fts else parse_bates_query(query_text) if incoming_bates[0] is not None and incoming_bates[1] is not None: parsed_bates = parse_bates_scope_input(query_text) if within: if "bates" not in scope and "keyword" in scope: raise RetrieverError("`/search --within` only composes within the current slot. Use `/search ` to add Bates alongside a keyword scope.") scope["bates"] = intersect_bates_scopes(scope.get("bates"), parsed_bates) else: scope["bates"] = parsed_bates return run_scope_search_from_session(root, paths, scope) if within: if "keyword" not in scope and "bates" in scope: raise RetrieverError("`/search --within` only composes within the current slot. Use `/search ` or `/filter ` to add a keyword alongside Bates scope.") existing_keyword = normalize_inline_whitespace(str(scope.get("keyword") or "")) scope["keyword"] = query_text if not existing_keyword else f"({existing_keyword}) AND ({query_text})" else: scope["keyword"] = query_text return run_scope_search_from_session(root, paths, scope) if command_name == "bates": if not normalized_tail: return {"status": "ok", "bates": scope.get("bates")} if normalized_tail == "clear": scope.pop("bates", None) return run_scope_search_from_session(root, paths, scope) scope["bates"] = parse_bates_scope_input(normalized_tail) return run_scope_search_from_session(root, paths, scope) if command_name == "filter": if not normalized_tail: return {"status": "ok", "filter": normalize_inline_whitespace(str(scope.get("filter") or "")) or None} if normalized_tail == "clear": scope.pop("filter", None) return run_scope_search_from_session(root, paths, scope) compile_sql_filter_expression(connection, normalized_tail) existing_filter = normalize_inline_whitespace(str(scope.get("filter") or "")) scope["filter"] = normalized_tail if not existing_filter else f"({existing_filter}) AND ({normalized_tail})" return run_scope_search_from_session(root, paths, scope) if command_name == "dataset": if not normalized_tail: return {"status": "ok", "dataset": coerce_scope_dataset_entries(scope.get("dataset"))} dataset_args = shlex_split_slash_tail(normalized_tail) if dataset_args and dataset_args[0] == "list": if len(dataset_args) != 1: raise RetrieverError("Usage: /dataset list") return {"status": "ok", "datasets": list_dataset_summaries(connection, root=root)} if dataset_args and dataset_args[0] == "clear": scope.pop("dataset", None) return run_scope_search_from_session(root, paths, scope) if dataset_args and dataset_args[0] == "rename": if len(dataset_args) != 3: raise RetrieverError("Usage: /dataset rename ") renamed_payload = rename_dataset(root, dataset_args[1], dataset_args[2]) dataset_summary = renamed_payload["dataset"] dataset_id = int(dataset_summary["id"]) dataset_name = str(dataset_summary["dataset_name"]) refresh_dataset_name_refs_in_scope(scope, dataset_id, dataset_name) updated_session_state = read_session_state(paths) updated_session_state["scope"] = coerce_scope_payload(scope) persist_session_state(paths, updated_session_state) refresh_dataset_name_refs_in_saved_scopes(paths, dataset_id, dataset_name) return renamed_payload dataset_names = split_quoted_comma_values(normalized_tail) if not dataset_names: raise RetrieverError("Dataset selection cannot be empty.") scope["dataset"] = resolve_scope_dataset_selection(connection, dataset_names) return run_scope_search_from_session(root, paths, scope) if command_name == "from-run": if not normalized_tail: return {"status": "ok", "from_run_id": scope.get("from_run_id")} if normalized_tail == "clear": scope.pop("from_run_id", None) return run_scope_search_from_session(root, paths, scope) scope["from_run_id"] = resolve_scope_from_run_id(connection, normalized_tail) return run_scope_search_from_session(root, paths, scope) raise RetrieverError(f"Unknown slash command: /{command_name}") finally: connection.close() def catalog(root: Path) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) reconcile_custom_fields_registry(connection, repair=True) intrinsic = [ catalog_field_entry(field_name, field_type, source="builtin") for field_name, field_type in sorted(BUILTIN_FIELD_TYPES.items()) if field_name not in CATALOG_EXCLUDED_BUILTIN_FIELDS ] custom_rows = connection.execute( """ SELECT field_name, field_type, instruction FROM custom_fields_registry ORDER BY field_name ASC """ ).fetchall() columns = table_columns(connection, "documents") custom = [ catalog_field_entry(str(row["field_name"]), str(row["field_type"]), source="custom", instruction=row["instruction"]) for row in custom_rows if row["field_name"] in columns and row["field_name"] not in CATALOG_EXCLUDED_CUSTOM_FIELDS ] virtual = [ catalog_field_entry(field_name, field_type, source="virtual") for field_name, field_type in sorted(VIRTUAL_FILTER_FIELD_TYPES.items()) ] return { "status": "ok", "intrinsic": intrinsic, "custom": custom, "virtual": virtual, } finally: connection.close() def normalize_entity_inventory_roles(raw_roles: list[str] | None) -> list[str]: if not raw_roles: return sorted(DOCUMENT_ENTITY_ROLES) roles: list[str] = [] for raw_role in raw_roles: normalized = normalize_entity_lookup_text(raw_role).replace(" ", "_") role = ENTITY_ROLE_FILTER_FIELDS.get(normalized, normalized) if role not in DOCUMENT_ENTITY_ROLES: raise RetrieverError( f"Unsupported entity role: {raw_role}. Supported roles: {', '.join(sorted(DOCUMENT_ENTITY_ROLES))}." ) if role not in roles: roles.append(role) return roles def scoped_entity_inventory_document_filter( connection: sqlite3.Connection, *, query: str = "", raw_filters: object | None = None, document_ids: list[int] | None = None, dataset_id: int | None = None, dataset_names: list[str] | None = None, conversation_id: int | None = None, from_run_id: int | None = None, ) -> tuple[list[str], list[object], list[object], dict[str, object]]: normalized_query = normalize_inline_whitespace(query) normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids or [])) normalized_dataset_names = [ name for raw_name in dataset_names or [] for name in [normalize_inline_whitespace(str(raw_name or ""))] if name ] scope: dict[str, object] = {} filter_summary: list[object] = [] if normalized_query: selection = resolve_document_search(connection, normalized_query, raw_filters, None, None) selected_ids = [int(item["id"]) for item in selection["results"]] filter_summary = list(selection["filters"]) if not selected_ids: return ["0"], [], filter_summary, {"query": normalized_query, "matched_document_count": 0} placeholders = ", ".join("?" for _ in selected_ids) clauses = [f"d.id IN ({placeholders})"] params: list[object] = selected_ids scope["query"] = normalized_query scope["matched_document_count"] = len(selected_ids) else: filter_summary, clauses, params = build_search_filters(connection, raw_filters) if normalized_document_ids: placeholders = ", ".join("?" for _ in normalized_document_ids) clauses.append(f"d.id IN ({placeholders})") params.extend(normalized_document_ids) scope["document_ids"] = normalized_document_ids if dataset_id is not None: clauses.append( """ EXISTS ( SELECT 1 FROM dataset_documents dd WHERE dd.document_id = d.id AND dd.dataset_id = ? ) """ ) params.append(int(dataset_id)) scope["dataset_id"] = int(dataset_id) if normalized_dataset_names: placeholders = ", ".join("?" for _ in normalized_dataset_names) clauses.append( f""" EXISTS ( SELECT 1 FROM dataset_documents dd JOIN datasets ds ON ds.id = dd.dataset_id WHERE dd.document_id = d.id AND ds.dataset_name IN ({placeholders}) ) """ ) params.extend(normalized_dataset_names) scope["dataset_names"] = normalized_dataset_names if conversation_id is not None: clauses.append("d.conversation_id = ?") params.append(int(conversation_id)) scope["conversation_id"] = int(conversation_id) if from_run_id is not None: run_id = resolve_scope_from_run_id(connection, from_run_id) clauses.append( """ EXISTS ( SELECT 1 FROM run_snapshot_documents rsd WHERE rsd.document_id = d.id AND rsd.run_id = ? ) """ ) params.append(run_id) scope["from_run_id"] = run_id if filter_summary: scope["filters"] = filter_summary return clauses, params, filter_summary, scope def entity_inventory_examples( connection: sqlite3.Connection, *, where_clause: str, scope_params: list[object], entity_id: int, role: str, limit: int, ) -> list[dict[str, object]]: if limit <= 0: return [] rows = connection.execute( f""" SELECT d.id, d.control_number, d.rel_path, d.title, d.date_created FROM documents d JOIN document_entities de ON de.document_id = d.id JOIN entities e ON e.id = de.entity_id WHERE {where_clause} AND de.entity_id = ? AND de.role = ? AND e.canonical_status = ? ORDER BY d.date_created DESC, d.id DESC LIMIT ? """, [*scope_params, entity_id, role, ENTITY_STATUS_ACTIVE, limit], ).fetchall() return [ { "document_id": int(row["id"]), "control_number": row["control_number"], "rel_path": row["rel_path"], "title": row["title"], "date_created": row["date_created"], } for row in rows ] def list_entity_role_inventory( root: Path, *, roles: list[str] | None = None, query: str = "", raw_filters: object | None = None, document_ids: list[int] | None = None, dataset_id: int | None = None, dataset_names: list[str] | None = None, conversation_id: int | None = None, from_run_id: int | None = None, limit: int = 100, examples_per_entity: int = 3, ) -> dict[str, object]: normalized_roles = normalize_entity_inventory_roles(roles) normalized_limit = max(1, min(int(limit or 100), 500)) normalized_examples_per_entity = max(0, min(int(examples_per_entity or 0), 10)) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) where_clauses, scope_params, filter_summary, scope = scoped_entity_inventory_document_filter( connection, query=query, raw_filters=raw_filters, document_ids=document_ids, dataset_id=dataset_id, dataset_names=dataset_names, conversation_id=conversation_id, from_run_id=from_run_id, ) where_clause = " AND ".join(where_clauses) role_placeholders = ", ".join("?" for _ in normalized_roles) rows = connection.execute( f""" SELECT de.role AS role, de.role AS roles, e.*, COUNT(DISTINCT de.document_id) AS document_count FROM documents d JOIN document_entities de ON de.document_id = d.id JOIN entities e ON e.id = de.entity_id WHERE {where_clause} AND de.role IN ({role_placeholders}) AND e.canonical_status = ? GROUP BY de.role, e.id ORDER BY de.role ASC, document_count DESC, COALESCE(e.sort_name, e.display_name, e.primary_email, e.primary_phone, '') ASC, e.id ASC LIMIT ? """, [*scope_params, *normalized_roles, ENTITY_STATUS_ACTIVE, normalized_limit], ).fetchall() entity_ids = [int(row["id"]) for row in rows] identifiers_by_entity = entity_identifiers_by_entity_id(connection, entity_ids) inventory_rows: list[dict[str, object]] = [] for row in rows: entity_id = int(row["id"]) role = str(row["role"]) identifiers = identifiers_by_entity.get(entity_id, []) summary = serialize_entity_summary(row, identifiers) inventory_rows.append( { "role": role, "entity_id": entity_id, "label": summary["label"], "entity_origin": summary["entity_origin"], "entity_type": summary["entity_type"], "primary_email": summary["primary_email"], "primary_phone": summary["primary_phone"], "document_count": summary["document_count"], "emails": summary["emails"], "phones": summary["phones"], "names": summary["names"], "examples": entity_inventory_examples( connection, where_clause=where_clause, scope_params=scope_params, entity_id=entity_id, role=role, limit=normalized_examples_per_entity, ), } ) return { "status": "ok", "roles": normalized_roles, "limit": normalized_limit, "examples_per_entity": normalized_examples_per_entity, "scope": scope, "filters": filter_summary, "rows": inventory_rows, } finally: connection.close() def fetch_visible_document_rows_by_ids( connection: sqlite3.Connection, document_ids: list[int], ) -> list[sqlite3.Row]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in document_ids)) if not normalized_document_ids: return [] placeholders = ", ".join("?" for _ in normalized_document_ids) rows = connection.execute( f""" SELECT d.*, EXISTS ( SELECT 1 FROM dataset_documents dd WHERE dd.document_id = d.id ) AS has_dataset_membership FROM documents d WHERE id IN ({placeholders}) """, normalized_document_ids, ).fetchall() rows_by_id = {int(row["id"]): row for row in rows} missing_ids: list[int] = [] lifecycle_hidden: list[str] = [] membership_hidden: list[int] = [] visible_rows_by_id: dict[int, sqlite3.Row] = {} for document_id in normalized_document_ids: row = rows_by_id.get(document_id) if row is None: missing_ids.append(document_id) continue if row["lifecycle_status"] in {"missing", "deleted"}: lifecycle_hidden.append(f"{document_id} ({row['lifecycle_status']})") continue if not bool(row["has_dataset_membership"]): membership_hidden.append(document_id) continue visible_rows_by_id[document_id] = row errors: list[str] = [] if missing_ids: errors.append( "Unknown document id" + ("" if len(missing_ids) == 1 else "s") + f": {', '.join(str(document_id) for document_id in missing_ids)}" ) if lifecycle_hidden: errors.append( "Document id" + ("" if len(lifecycle_hidden) == 1 else "s") + " not visible due to lifecycle_status: " + ", ".join(lifecycle_hidden) ) if membership_hidden: errors.append( "Document id" + ("" if len(membership_hidden) == 1 else "s") + " not visible because they have no dataset memberships: " + ", ".join(str(document_id) for document_id in membership_hidden) ) if errors: raise RetrieverError(" ".join(errors)) return [visible_rows_by_id[document_id] for document_id in normalized_document_ids] def resolve_fill_document_refs(connection: sqlite3.Connection, raw_doc_refs: list[str]) -> list[int]: resolved_document_ids: list[int] = [] for raw_doc_ref in raw_doc_refs: doc_ref = normalize_inline_whitespace(str(raw_doc_ref or "")) if not doc_ref: continue if doc_ref.isdigit(): resolved_document_ids.append(int(doc_ref)) continue rows = connection.execute( """ SELECT id FROM documents WHERE control_number = ? ORDER BY id ASC """, (doc_ref,), ).fetchall() if not rows: raise RetrieverError(f"Unknown document reference: {doc_ref}") if len(rows) > 1: raise RetrieverError(f"Document reference '{doc_ref}' is ambiguous.") resolved_document_ids.append(int(rows[0]["id"])) fetch_visible_document_rows_by_ids(connection, resolved_document_ids) return list(dict.fromkeys(resolved_document_ids)) def fetch_attachment_parent_ids( connection: sqlite3.Connection, parent_ids: list[int], ) -> set[int]: if not parent_ids: return set() placeholders = ", ".join("?" for _ in parent_ids) rows = connection.execute( f""" SELECT DISTINCT parent_document_id FROM documents WHERE parent_document_id IN ({placeholders}) AND COALESCE(child_document_kind, ?) = ? AND lifecycle_status NOT IN ('missing', 'deleted') """, [*parent_ids, CHILD_DOCUMENT_KIND_ATTACHMENT, CHILD_DOCUMENT_KIND_ATTACHMENT], ).fetchall() return { int(row["parent_document_id"]) for row in rows if row["parent_document_id"] is not None } def resolve_export_field_definitions( connection: sqlite3.Connection, raw_fields: list[str] | None, ) -> list[dict[str, str]]: if not raw_fields: raise RetrieverError("export-csv requires at least one --field.") resolved_fields: list[dict[str, str]] = [] seen_fields: set[str] = set() for raw_field in raw_fields: field_def = resolve_field_definition(connection, raw_field) normalized_field_name = str(field_def["field_name"]) if normalized_field_name in seen_fields: raise RetrieverError(f"Duplicate export field: {normalized_field_name}") seen_fields.add(normalized_field_name) resolved_fields.append( { "requested_name": raw_field, "field_name": normalized_field_name, "field_type": str(field_def["field_type"]), "source": str(field_def.get("source") or ""), } ) return resolved_fields ARCHIVE_MANIFEST_REL_PATH = ".retriever/export-manifest.json" ARCHIVE_ROOT_MANIFEST_REL_PATH = "manifest.json" ARCHIVE_METADATA_REL_PATH = "metadata.csv" ARCHIVE_ROOT_README_REL_PATH = "README.txt" ARCHIVE_PREVIEW_ROOT_REL_PATH = "previews" ARCHIVE_PREVIEWS_INDEX_REL_PATH = "previews.csv" ARCHIVE_SOURCE_PARTS_REL_PATH = "source_parts.csv" ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH = "extracted_text.csv" ARCHIVE_TEXT_ROOT_REL_PATH = "text" ARCHIVE_CHECKSUMS_REL_PATH = "checksums.csv" ARCHIVE_LOADFILE_REL_PATH = "loadfile.dat" ARCHIVE_METADATA_BASE_FIELD_NAMES = ( "control_number", "title", "subject", "file_name", "content_type", "author", "recipients", "participants", "date_created", "date_modified", "dataset_name", "production_name", "custodian", "page_count", "rel_path", "source_kind", "source_rel_path", "source_folder_path", "conversation_id", "parent_document_id", "child_document_kind", "is_attachment", "has_attachments", "begin_bates", "end_bates", "begin_attachment", "end_attachment", "updated_at", ) ARCHIVE_METADATA_EXTRA_FIELD_DEFS = ( {"field_name": "document_id", "field_type": "integer", "source": "archive"}, {"field_name": "document_entry_kind", "field_type": "text", "source": "archive"}, {"field_name": "preview_primary_rel_path", "field_type": "text", "source": "archive"}, {"field_name": "preview_rel_paths", "field_type": "text", "source": "archive"}, {"field_name": "native_source_rel_paths", "field_type": "text", "source": "archive"}, {"field_name": "text_source_rel_paths", "field_type": "text", "source": "archive"}, {"field_name": "image_source_rel_paths", "field_type": "text", "source": "archive"}, {"field_name": "source_part_rel_paths", "field_type": "text", "source": "archive"}, ) def archive_metadata_custom_field_names(connection: sqlite3.Connection) -> list[str]: if not table_exists(connection, "custom_fields_registry"): return [] document_columns = set(table_columns(connection, "documents")) names: list[str] = [] seen: set[str] = set() for row in connection.execute( """ SELECT field_name FROM custom_fields_registry ORDER BY field_name ASC """ ).fetchall(): field_name = normalize_inline_whitespace(str(row["field_name"] or "")) if not field_name or field_name in seen or field_name not in document_columns: continue seen.add(field_name) names.append(field_name) return names def resolve_archive_metadata_field_definitions(connection: sqlite3.Connection) -> list[dict[str, str]]: base_field_defs = resolve_export_field_definitions( connection, [*ARCHIVE_METADATA_BASE_FIELD_NAMES, *archive_metadata_custom_field_names(connection)], ) return [ dict(ARCHIVE_METADATA_EXTRA_FIELD_DEFS[0]), *base_field_defs, *[dict(field_def) for field_def in ARCHIVE_METADATA_EXTRA_FIELD_DEFS[1:]], ] def archive_source_part_archive_rel_path(source_part_entry: dict[str, object]) -> str: return normalize_whitespace( str(source_part_entry.get("archive_rel_path") or source_part_entry.get("rel_path") or "") ) def archive_metadata_field_value( row: sqlite3.Row, field_def: dict[str, str], context: dict[str, object], manifest_entry: dict[str, object], ) -> object: field_name = str(field_def["field_name"]) if str(field_def.get("source") or "") != "archive": return export_field_value(row, field_def, context) preview_rel_paths = [ normalize_whitespace(str(path or "")) for path in list(manifest_entry.get("preview_rel_paths") or []) if normalize_whitespace(str(path or "")) ] source_part_entries = [ dict(entry) for entry in list(manifest_entry.get("source_part_entries") or []) if isinstance(entry, dict) ] source_part_rel_paths = [ rel_path for rel_path in (archive_source_part_archive_rel_path(entry) for entry in source_part_entries) if rel_path ] native_source_rel_paths = [ rel_path for rel_path in ( archive_source_part_archive_rel_path(entry) for entry in source_part_entries if normalize_whitespace(str(entry.get("part_kind") or "")).lower() == "native" ) if rel_path ] text_source_rel_paths = [ rel_path for rel_path in ( archive_source_part_archive_rel_path(entry) for entry in source_part_entries if normalize_whitespace(str(entry.get("part_kind") or "")).lower() == "text" ) if rel_path ] image_source_rel_paths = [ rel_path for rel_path in ( archive_source_part_archive_rel_path(entry) for entry in source_part_entries if normalize_whitespace(str(entry.get("part_kind") or "")).lower() == "image" ) if rel_path ] if field_name == "document_id": return int(manifest_entry.get("document_id") or row["id"]) if field_name == "document_entry_kind": return manifest_entry.get("document_entry_kind") if field_name == "preview_primary_rel_path": return preview_rel_paths[0] if preview_rel_paths else None if field_name == "preview_rel_paths": return preview_rel_paths if field_name == "native_source_rel_paths": return native_source_rel_paths if field_name == "text_source_rel_paths": return text_source_rel_paths if field_name == "image_source_rel_paths": return image_source_rel_paths if field_name == "source_part_rel_paths": return source_part_rel_paths raise RetrieverError(f"Unknown archive metadata field: {field_name}") def build_archive_metadata_csv_bytes( connection: sqlite3.Connection, rows: list[sqlite3.Row], manifest_entries: list[dict[str, object]], field_defs: list[dict[str, str]], ) -> bytes: rows_by_id = {int(row["id"]): row for row in rows} context = build_export_context(connection, rows, field_defs) buffer = io.StringIO() writer = csv.writer(buffer) writer.writerow([field_def["field_name"] for field_def in field_defs]) for manifest_entry in manifest_entries: try: document_id = int(manifest_entry["document_id"]) except (KeyError, TypeError, ValueError): continue row = rows_by_id.get(document_id) if row is None: continue writer.writerow( [ serialize_export_cell_value( archive_metadata_field_value(row, field_def, context, manifest_entry), str(field_def["field_type"]), ) for field_def in field_defs ] ) return buffer.getvalue().encode("utf-8") def serialize_archive_csv_cell_value(value: object) -> str: if value is None: return "" if isinstance(value, bool): return "true" if value else "false" if isinstance(value, (list, tuple)): return "; ".join(str(item) for item in value if item is not None) return str(value) def build_archive_csv_bytes(field_names: list[str], row_dicts: list[dict[str, object]]) -> bytes: buffer = io.StringIO() writer = csv.writer(buffer) writer.writerow(field_names) for row_dict in row_dicts: writer.writerow([serialize_archive_csv_cell_value(row_dict.get(field_name)) for field_name in field_names]) return buffer.getvalue().encode("utf-8") def archive_safe_token(raw_value: object, *, default: str) -> str: normalized = normalize_whitespace(str(raw_value or "")) if not normalized: return default candidate = re.sub(r"[^A-Za-z0-9._-]+", "-", normalized).strip("-._") return candidate or default def archive_windows_loadfile_path(raw_rel_path: str) -> str: return ".\\" + normalize_archive_member_path(raw_rel_path).replace("/", "\\") def archive_text_document_rel_path(document_id: int, control_number: object) -> str: default_stem = f"document-{int(document_id):08d}" stem = archive_safe_token(control_number, default=default_stem) if stem == default_stem: file_name = f"{stem}.txt" else: file_name = f"{stem}-{int(document_id):08d}.txt" return normalize_archive_member_path(str(Path(ARCHIVE_TEXT_ROOT_REL_PATH) / file_name)) def build_archive_dat_line(fields: list[object]) -> bytes: delimiter = b"\x14" quote = b"\xfe" return delimiter.join(quote + str(field or "").encode("utf-8") + quote for field in fields) + b"\r\n" def prepare_archive_sidecar_entries( connection: sqlite3.Connection, paths: dict[str, Path], rows: list[sqlite3.Row], manifest_entries: list[dict[str, object]], ) -> dict[str, object]: rows_by_id = {int(row["id"]): row for row in rows} preview_index_rows: list[dict[str, object]] = [] source_part_index_rows: list[dict[str, object]] = [] extracted_text_index_rows: list[dict[str, object]] = [] text_file_entries: list[dict[str, object]] = [] loadfile_rows: list[list[object]] = [] counts = { "documents_with_text": 0, "missing_preview_entries": 0, "missing_source_part_entries": 0, "preview_entries": 0, "source_part_entries": 0, "synthetic_documents": 0, } for manifest_entry in manifest_entries: try: document_id = int(manifest_entry["document_id"]) except (KeyError, TypeError, ValueError): continue row = rows_by_id.get(document_id) if row is None: continue control_number = row["control_number"] document_rel_path = normalize_archive_member_path( str(manifest_entry.get("rel_path") or row["rel_path"]), label=f"Archive rel_path for document {document_id}", ) document_entry_kind = normalize_whitespace(str(manifest_entry.get("document_entry_kind") or "")) if document_entry_kind == "synthetic": counts["synthetic_documents"] += 1 exported_preview_paths = { normalize_archive_member_path(str(path), label=f"Preview path for document {document_id}") for path in list(manifest_entry.get("preview_rel_paths") or []) if normalize_whitespace(str(path or "")) } for preview_row in document_preview_rows(connection, document_id): preview_rel_path = normalize_archive_member_path( str(preview_row["rel_preview_path"]), label=f"Preview path for document {document_id}", ) is_missing = preview_rel_path not in exported_preview_paths preview_index_rows.append( { "document_id": document_id, "control_number": control_number, "document_rel_path": document_rel_path, "preview_type": preview_row["preview_type"], "ordinal": int(preview_row["ordinal"] or 0), "label": preview_row["label"], "rel_path": preview_rel_path, "is_missing": is_missing, } ) counts["preview_entries"] += 1 if is_missing: counts["missing_preview_entries"] += 1 native_rel_path = "" for source_part_entry in list(manifest_entry.get("source_part_entries") or []): if not isinstance(source_part_entry, dict): continue rel_path = archive_source_part_archive_rel_path(source_part_entry) if not rel_path: continue part_kind = normalize_whitespace(str(source_part_entry.get("part_kind") or "")) is_missing = bool(source_part_entry.get("missing")) source_part_index_rows.append( { "document_id": document_id, "control_number": control_number, "document_rel_path": document_rel_path, "part_kind": part_kind, "ordinal": int(source_part_entry.get("ordinal") or 0), "label": source_part_entry.get("label"), "rel_path": rel_path, "is_missing": is_missing, } ) counts["source_part_entries"] += 1 if is_missing: counts["missing_source_part_entries"] += 1 continue if part_kind.lower() == "native" and not native_rel_path: native_rel_path = rel_path text_content = document_source_text_body(connection, paths, row) text_rel_path = "" text_char_count = 0 text_line_count = 0 if text_content: text_rel_path = archive_text_document_rel_path(document_id, control_number) text_char_count = len(text_content) text_line_count = text_content.count("\n") + 1 text_file_entries.append( { "document_id": document_id, "control_number": control_number, "rel_path": text_rel_path, "payload_bytes": text_content.encode("utf-8"), } ) counts["documents_with_text"] += 1 manifest_entry["text_rel_path"] = text_rel_path else: manifest_entry["text_rel_path"] = None extracted_text_index_rows.append( { "document_id": document_id, "control_number": control_number, "document_rel_path": document_rel_path, "text_rel_path": text_rel_path, "char_count": text_char_count, "line_count": text_line_count, "has_text": bool(text_rel_path), } ) primary_document_rel_path = document_rel_path if document_entry_kind == "copied" else "" if not native_rel_path: native_rel_path = primary_document_rel_path loadfile_rows.append( [ row["begin_bates"] or row["control_number"] or "", row["end_bates"] or row["control_number"] or "", row["begin_attachment"] or "", row["end_attachment"] or "", document_id, row["control_number"] or "", row["file_name"] or "", row["title"] or "", row["subject"] or "", row["author"] or "", row["date_created"] or "", row["date_modified"] or "", document_custodian_display_text_from_row(row), archive_windows_loadfile_path(text_rel_path) if text_rel_path else "", archive_windows_loadfile_path(native_rel_path) if native_rel_path else "", archive_windows_loadfile_path(primary_document_rel_path) if primary_document_rel_path else "", ] ) source_parts_payload = build_archive_csv_bytes( [ "document_id", "control_number", "document_rel_path", "part_kind", "ordinal", "label", "rel_path", "is_missing", ], source_part_index_rows, ) previews_payload = build_archive_csv_bytes( [ "document_id", "control_number", "document_rel_path", "preview_type", "ordinal", "label", "rel_path", "is_missing", ], preview_index_rows, ) extracted_text_payload = build_archive_csv_bytes( [ "document_id", "control_number", "document_rel_path", "text_rel_path", "char_count", "line_count", "has_text", ], extracted_text_index_rows, ) loadfile_headers = [ "BEGBATES", "ENDBATES", "BEGATTACH", "ENDATTACH", "DOCID", "CONTROL_NUMBER", "FILE_NAME", "TITLE", "SUBJECT", "AUTHOR", "DATE_CREATED", "DATE_MODIFIED", "CUSTODIAN", "TEXT_PATH", "NATIVE_PATH", "DOCUMENT_PATH", ] loadfile_payload = build_archive_dat_line(loadfile_headers) + b"".join( build_archive_dat_line(loadfile_row) for loadfile_row in loadfile_rows ) return { "counts": counts, "extracted_text_index_payload": extracted_text_payload, "loadfile_payload": loadfile_payload, "preview_index_payload": previews_payload, "source_part_index_payload": source_parts_payload, "text_file_entries": text_file_entries, } def build_archive_root_manifest_payload( *, created_at: str, selector: dict[str, object], family_mode: str | None, seed_limit: int | None, portable_workspace: bool, portable_workspace_payload: dict[str, object] | None, document_count: int, failed_count: int, warning_count: int, sidecar_counts: dict[str, object], ) -> dict[str, object]: return { "format": "retriever-archive", "version": 2, "status": "ok" if failed_count == 0 else "failed", "created_at": created_at, "selector": selector, "family_mode": family_mode, "seed_limit": seed_limit, "portable_workspace": portable_workspace, "files": { "metadata_rel_path": ARCHIVE_METADATA_REL_PATH, "preview_root_rel_path": ARCHIVE_PREVIEW_ROOT_REL_PATH, "previews_index_rel_path": ARCHIVE_PREVIEWS_INDEX_REL_PATH, "source_parts_rel_path": ARCHIVE_SOURCE_PARTS_REL_PATH, "extracted_text_index_rel_path": ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH, "text_root_rel_path": ARCHIVE_TEXT_ROOT_REL_PATH, "loadfile_rel_path": ARCHIVE_LOADFILE_REL_PATH, "checksums_rel_path": ARCHIVE_CHECKSUMS_REL_PATH, "readme_rel_path": ARCHIVE_ROOT_README_REL_PATH, "internal_manifest_rel_path": ARCHIVE_MANIFEST_REL_PATH, "portable_workspace_db_rel_path": ".retriever/retriever.db" if portable_workspace else None, }, "counts": { "document_count": document_count, "documents_with_text": int(sidecar_counts.get("documents_with_text") or 0), "preview_entries": int(sidecar_counts.get("preview_entries") or 0), "source_part_entries": int(sidecar_counts.get("source_part_entries") or 0), "synthetic_documents": int(sidecar_counts.get("synthetic_documents") or 0), "warning_count": warning_count, "failed_item_count": failed_count, }, "portable_workspace_document_ids": ( portable_workspace_payload.get("selected_document_ids") if isinstance(portable_workspace_payload, dict) else [] ), "portable_workspace_stub_document_ids": ( portable_workspace_payload.get("stub_document_ids") if isinstance(portable_workspace_payload, dict) else [] ), } def build_archive_readme_bytes(root_manifest_payload: dict[str, object]) -> bytes: counts = root_manifest_payload.get("counts") if isinstance(root_manifest_payload, dict) else {} if not isinstance(counts, dict): counts = {} lines = [ "Retriever Archive Export", "", f"Created: {root_manifest_payload.get('created_at') or ''}", f"Documents: {int(counts.get('document_count') or 0)}", f"Warnings: {int(counts.get('warning_count') or 0)}", "", "Contents:", f"- {ARCHIVE_METADATA_REL_PATH}: one row per logical document with core metadata and archive paths.", f"- {ARCHIVE_ROOT_MANIFEST_REL_PATH}: importer-facing archive summary and sidecar locations.", f"- {ARCHIVE_ROOT_README_REL_PATH}: this guide.", f"- {ARCHIVE_PREVIEW_ROOT_REL_PATH}/: exported preview artifacts for human review.", f"- {ARCHIVE_PREVIEWS_INDEX_REL_PATH}: normalized preview rows with preview types and paths.", f"- {ARCHIVE_SOURCE_PARTS_REL_PATH}: normalized source/native/image part rows.", f"- {ARCHIVE_TEXT_ROOT_REL_PATH}/: UTF-8 extracted text files when text is available.", f"- {ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH}: per-document extracted text paths and counts.", f"- {ARCHIVE_LOADFILE_REL_PATH}: Concordance-style DAT with text/native/document paths.", f"- {ARCHIVE_CHECKSUMS_REL_PATH}: SHA-256 and byte size for every archive member except {ARCHIVE_CHECKSUMS_REL_PATH}.", f"- {ARCHIVE_MANIFEST_REL_PATH}: Retriever-specific lossless manifest with full document details.", "", "Notes:", "- Source and native artifacts stay at the same relative paths they had when Retriever ingested them.", "- Synthetic .logical entries may appear when one review document is composed from multiple underlying parts.", "- If a portable workspace was requested, .retriever/retriever.db contains a searchable metadata subset for the selected documents.", ] return ("\n".join(lines) + "\n").encode("utf-8") def write_archive_sidecars( archive: zipfile.ZipFile, written_member_paths: set[str], connection: sqlite3.Connection, paths: dict[str, Path], rows: list[sqlite3.Row], manifest_document_entries: list[dict[str, object]], *, metadata_payload: bytes, created_at: str, selector: dict[str, object], family_mode: str | None, seed_limit: int | None, portable_workspace: bool, portable_workspace_payload: dict[str, object] | None, warnings: list[str], failed_count: int, ) -> dict[str, object]: add_archive_bytes_once( archive, written_member_paths, metadata_payload, ARCHIVE_METADATA_REL_PATH, ) sidecar_state = prepare_archive_sidecar_entries( connection, paths, rows, manifest_document_entries, ) for text_file_entry in sidecar_state["text_file_entries"]: add_archive_bytes_once( archive, written_member_paths, bytes(text_file_entry["payload_bytes"]), str(text_file_entry["rel_path"]), ) add_archive_bytes_once( archive, written_member_paths, sidecar_state["preview_index_payload"], ARCHIVE_PREVIEWS_INDEX_REL_PATH, ) add_archive_bytes_once( archive, written_member_paths, sidecar_state["source_part_index_payload"], ARCHIVE_SOURCE_PARTS_REL_PATH, ) add_archive_bytes_once( archive, written_member_paths, sidecar_state["extracted_text_index_payload"], ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH, ) add_archive_bytes_once( archive, written_member_paths, sidecar_state["loadfile_payload"], ARCHIVE_LOADFILE_REL_PATH, ) root_manifest_payload = build_archive_root_manifest_payload( created_at=created_at, selector=selector, family_mode=family_mode, seed_limit=seed_limit, portable_workspace=portable_workspace, portable_workspace_payload=portable_workspace_payload, document_count=len(manifest_document_entries), failed_count=failed_count, warning_count=len(warnings), sidecar_counts=dict(sidecar_state["counts"]), ) add_archive_bytes_once( archive, written_member_paths, (json.dumps(root_manifest_payload, indent=2, sort_keys=True) + "\n").encode("utf-8"), ARCHIVE_ROOT_MANIFEST_REL_PATH, ) add_archive_bytes_once( archive, written_member_paths, build_archive_readme_bytes(root_manifest_payload), ARCHIVE_ROOT_README_REL_PATH, ) internal_manifest_payload = { "status": "ok" if failed_count == 0 else "failed", "created_at": created_at, "selector": selector, "family_mode": family_mode, "seed_limit": seed_limit, "document_count": len(manifest_document_entries), "portable_workspace": portable_workspace, "metadata_rel_path": ARCHIVE_METADATA_REL_PATH, "preview_root_rel_path": ARCHIVE_PREVIEW_ROOT_REL_PATH, "root_manifest_rel_path": ARCHIVE_ROOT_MANIFEST_REL_PATH, "readme_rel_path": ARCHIVE_ROOT_README_REL_PATH, "previews_index_rel_path": ARCHIVE_PREVIEWS_INDEX_REL_PATH, "source_parts_rel_path": ARCHIVE_SOURCE_PARTS_REL_PATH, "extracted_text_index_rel_path": ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH, "text_root_rel_path": ARCHIVE_TEXT_ROOT_REL_PATH, "checksums_rel_path": ARCHIVE_CHECKSUMS_REL_PATH, "loadfile_rel_path": ARCHIVE_LOADFILE_REL_PATH, "portable_workspace_document_ids": ( portable_workspace_payload["selected_document_ids"] if portable_workspace_payload is not None else [] ), "portable_workspace_stub_document_ids": ( portable_workspace_payload["stub_document_ids"] if portable_workspace_payload is not None else [] ), "counts": dict(sidecar_state["counts"]), "documents": manifest_document_entries, "warnings": warnings, "failed_items": failed_count, } add_archive_bytes_once( archive, written_member_paths, (json.dumps(internal_manifest_payload, indent=2, sort_keys=True) + "\n").encode("utf-8"), ARCHIVE_MANIFEST_REL_PATH, ) return { "internal_manifest_payload": internal_manifest_payload, "root_manifest_payload": root_manifest_payload, "sidecar_counts": dict(sidecar_state["counts"]), } def build_archive_checksums_csv_bytes(archive: zipfile.ZipFile, member_paths: list[str]) -> bytes: buffer = io.StringIO() writer = csv.writer(buffer) writer.writerow(["rel_path", "file_size_bytes", "sha256"]) for member_path in member_paths: info = archive.getinfo(member_path) digest = hashlib.sha256() with archive.open(member_path, "r") as handle: while True: chunk = handle.read(1024 * 1024) if not chunk: break digest.update(chunk) writer.writerow([member_path, int(info.file_size), digest.hexdigest()]) return buffer.getvalue().encode("utf-8") def ensure_archive_checksums_file(archive_path: Path) -> dict[str, int]: with zipfile.ZipFile(archive_path, "a", compression=zipfile.ZIP_DEFLATED) as archive: member_paths = sorted(set(archive.namelist())) if ARCHIVE_CHECKSUMS_REL_PATH not in member_paths: checksummed_member_paths = [member_path for member_path in member_paths if member_path != ARCHIVE_CHECKSUMS_REL_PATH] archive.writestr( ARCHIVE_CHECKSUMS_REL_PATH, build_archive_checksums_csv_bytes(archive, checksummed_member_paths), ) member_paths = checksummed_member_paths + [ARCHIVE_CHECKSUMS_REL_PATH] checksummed_count = len([member_path for member_path in member_paths if member_path != ARCHIVE_CHECKSUMS_REL_PATH]) return { "archive_member_count": len(member_paths), "checksummed_member_count": checksummed_count, } def resolve_table_export_field_definitions( connection: sqlite3.Connection, table_name: str, raw_fields: list[str] | None, ) -> list[dict[str, str]]: normalized_table_name = normalize_export_table_name(table_name) if normalized_table_name == "documents": return resolve_export_field_definitions(connection, raw_fields or default_display_columns(BROWSE_MODE_DOCUMENTS)) fields = list(raw_fields or default_display_columns( BROWSE_MODE_ENTITIES if normalized_table_name == "entities" else BROWSE_MODE_CONVERSATIONS )) if not fields: raise RetrieverError(f"/export table {normalized_table_name} requires at least one field.") browse_mode = BROWSE_MODE_ENTITIES if normalized_table_name == "entities" else BROWSE_MODE_CONVERSATIONS resolved_fields: list[dict[str, str]] = [] seen_fields: set[str] = set() for raw_field in fields: field_def = resolve_browse_field_definition(connection, raw_field, browse_mode=browse_mode) normalized_field_name = str(field_def["field_name"]) if normalized_field_name in seen_fields: raise RetrieverError(f"Duplicate export field: {normalized_field_name}") seen_fields.add(normalized_field_name) resolved_fields.append( { "requested_name": raw_field, "field_name": normalized_field_name, "field_type": str(field_def["field_type"]), "source": str(field_def.get("source") or ""), } ) return resolved_fields def path_within(candidate: Path, parent: Path) -> bool: try: candidate.relative_to(parent) return True except ValueError: return False def resolve_export_output_path(paths: dict[str, Path], raw_output_path: str) -> Path: normalized_output_path = raw_output_path.strip() if not normalized_output_path: raise RetrieverError("Output path cannot be empty.") exports_dir = (paths["state_dir"] / "exports").resolve() requested_path = Path(normalized_output_path).expanduser() if requested_path.is_absolute(): resolved_path = requested_path.resolve() else: resolved_path = (exports_dir / requested_path).resolve() if not path_within(resolved_path, exports_dir): raise RetrieverError(f"Relative output paths must stay within {exports_dir}.") if resolved_path.exists() and resolved_path.is_dir(): raise RetrieverError(f"Output path is a directory: {resolved_path}") workspace_root = paths["root"].resolve() if path_within(resolved_path, workspace_root) and not path_within(resolved_path, exports_dir): raise RetrieverError( f"Workspace-internal output paths must live under {exports_dir} to avoid re-ingesting exported artifacts." ) return resolved_path def resolve_export_output_dir(paths: dict[str, Path], raw_output_path: str) -> Path: normalized_output_path = raw_output_path.strip() if not normalized_output_path: raise RetrieverError("Output path cannot be empty.") exports_dir = (paths["state_dir"] / "exports").resolve() requested_path = Path(normalized_output_path).expanduser() if requested_path.is_absolute(): resolved_path = requested_path.resolve() else: resolved_path = (exports_dir / requested_path).resolve() if not path_within(resolved_path, exports_dir): raise RetrieverError(f"Relative output paths must stay within {exports_dir}.") if resolved_path.exists() and not resolved_path.is_dir(): raise RetrieverError(f"Output path is not a directory: {resolved_path}") workspace_root = paths["root"].resolve() if path_within(resolved_path, workspace_root) and not path_within(resolved_path, exports_dir): raise RetrieverError( f"Workspace-internal output paths must live under {exports_dir} to avoid re-ingesting exported exports." ) return resolved_path def relative_output_path_or_none(root: Path, output_path: Path) -> str | None: try: return output_path.relative_to(root).as_posix() except ValueError: return None def export_preview_unit_file_name(unit: dict[str, object]) -> str: unit_kind = str(unit["unit_kind"]) if unit_kind == "email_conversation": return f"conversation-{int(unit['conversation_id']):08d}.html" if unit_kind == "conversation_run": return ( f"conversation-{int(unit['conversation_id']):08d}" f"-run-{int(unit['run_start_index']) + 1:04d}-{int(unit['run_end_index']) + 1:04d}.html" ) return f"document-{int(unit['documents'][0]['id']):08d}.html" def export_preview_document_file_name(document: dict[str, object]) -> str: return f"document-{int(document['id']):08d}.html" def copy_export_preview_assets( paths: dict[str, Path], output_dir: Path, documents: list[dict[str, object]], ) -> list[str]: copied_rel_paths: list[str] = [] seen_rel_paths: set[str] = set() for document in documents: body_html = document.get("standalone_preview_body_html") if not isinstance(body_html, str) or not body_html.strip(): continue for asset_rel_path in preview_relative_asset_rel_paths_for_html_body( body_html, source_preview_rel_path=document.get("standalone_preview_rel_path"), ): if asset_rel_path in seen_rel_paths: continue source_path = paths["state_dir"] / asset_rel_path if not source_path.exists() or not source_path.is_file(): continue target_path = output_dir / asset_rel_path target_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(source_path, target_path) seen_rel_paths.add(asset_rel_path) copied_rel_paths.append(asset_rel_path) return copied_rel_paths def build_export_preview_document_html( connection: sqlite3.Connection, paths: dict[str, Path], *, unit: dict[str, object], document: dict[str, object], document_output_path: Path, document_output_rel_path: str, unit_output_path: Path, ) -> str: document_id = int(document["id"]) document_heading = conversation_preview_document_heading(document) or f"Document {document_id}" content_type = normalize_whitespace(str(document.get("content_type") or "")).lower() unit_documents = list(unit["documents"]) if content_type == "email": conversation_row = None conversation_documents = None position_index = None thread_link_href = None if str(unit["unit_kind"]) == "email_conversation" and len(unit_documents) > 1: conversation_row = {"display_name": str(unit["title"])} conversation_documents = unit_documents email_document_ids = [ int(item["id"]) for item in unit_documents if normalize_whitespace(str(item.get("content_type") or "")).lower() == "email" ] if document_id in email_document_ids: position_index = email_document_ids.index(document_id) + 1 thread_link_href = relative_preview_href(unit_output_path, document_output_path) body_html = ( str(document.get("standalone_preview_body_html")) if isinstance(document.get("standalone_preview_body_html"), str) and str(document.get("standalone_preview_body_html")).strip() else None ) return build_email_message_preview_html( document, body_html=body_html, conversation_row=conversation_row, conversation_documents=conversation_documents, position_index=position_index, thread_link_href=thread_link_href, target_preview_rel_path=document_output_rel_path, ) attachment_links = ( conversation_attachment_links_by_document_id( connection, paths, segment_preview_path=document_output_path, documents=[document], ).get(document_id) or [] ) section_html = render_conversation_document_section( document, current_segment_href=document_output_path.name, doc_target_hrefs={document_id: f"#{conversation_preview_anchor(document_id)}"}, attachment_links_by_document_id={document_id: attachment_links} if attachment_links else None, target_preview_rel_path=document_output_rel_path, ) nav_links = ['Contents'] if len(unit_documents) > 1: nav_label = "Thread" if str(unit["unit_kind"]) == "email_conversation" else "Group" nav_links.append( f'{html.escape(nav_label)}' ) headers = { "Preview": document_heading, "Type": conversation_preview_document_kind(document), } if len(unit_documents) > 1: headers["Group"] = str(unit["title"]) return build_html_preview( headers, body_html=( "
    " '
    ' f'
    {"".join(nav_links)}
    ' "
    " f"{section_html}" "
    " ), document_title=document_heading, head_html=build_conversation_preview_head_html(), heading=document_heading, ) def build_export_preview_unit_html( unit: dict[str, object], *, file_name: str, output_rel_path: str, ) -> str: documents = list(unit["documents"]) selected_document_ids = {int(document_id) for document_id in unit["selected_document_ids"]} if unit["unit_kind"] == "standalone": heading = conversation_preview_document_heading(documents[0]) preview_type = "document" elif unit["unit_kind"] == "email_conversation": heading = str(unit["title"]) preview_type = "email conversation" else: heading = str(unit["title"]) preview_type = "conversation run" headers = { "Preview": heading, "Type": preview_type, "Documents": str(len(documents)), "Selected": str(len(selected_document_ids)), } if unit.get("conversation_type"): headers[passive_field_label("conversation_type", mixed_context=True)] = str(unit["conversation_type"]) doc_target_hrefs = { int(document["id"]): f"#{conversation_preview_anchor(int(document['id']))}" for document in documents } sections: list[str] = [] for document in documents: section_html = render_conversation_document_section( document, current_segment_href=file_name, doc_target_hrefs=doc_target_hrefs, target_preview_rel_path=output_rel_path, ) if int(document["id"]) in selected_document_ids: section_html = section_html.replace( '
    " "
    " f"{nav_links}" "
    " f"{''.join(sections)}" "" ), document_title=heading, head_html=build_conversation_preview_head_html(), heading=heading, ) def build_export_preview_index_html( *, units: list[dict[str, object]], selected_rows: list[sqlite3.Row], document_targets_by_id: dict[int, dict[str, object]], ) -> str: headers = { "Selected documents": str(len(selected_rows)), "Preview files": str(len(units)), } cards: list[str] = [] for unit in units: output_rel_path = str(unit["output_rel_path"]) selected_links = "".join( f"
  • {html.escape(str(document_targets_by_id[int(document_id)]['title']))}
    • " for document_id in unit["selected_document_ids"] if int(document_id) in document_targets_by_id ) cards.append( "
    " f"

    {html.escape(str(unit['title']))}

    " f"

    {html.escape(str(unit['summary']))}

    " f"{'
      ' + selected_links + '
    ' if selected_links else ''}" "
    " ) return build_html_preview( headers, body_html=f"
    {''.join(cards)}
    ", document_title="Export Preview Index", head_html=build_conversation_preview_head_html(), heading="Export Preview Index", ) def cleanup_previous_export_preview_outputs(output_dir: Path) -> None: manifest_path = output_dir / "manifest.json" if not manifest_path.exists(): return try: manifest_payload = json.loads(manifest_path.read_text(encoding="utf-8")) except Exception: return relative_paths: set[str] = set() for key in ("index_rel_path",): value = manifest_payload.get(key) if isinstance(value, str) and value: relative_paths.add(value) for unit in manifest_payload.get("units", []): if isinstance(unit, dict): value = unit.get("output_rel_path") if isinstance(value, str) and value: relative_paths.add(value) for document_target in manifest_payload.get("document_targets", []): if isinstance(document_target, dict): value = document_target.get("output_rel_path") if isinstance(value, str) and value: relative_paths.add(value) for value in manifest_payload.get("asset_rel_paths", []): if isinstance(value, str) and value: relative_paths.add(value) relative_paths.add("manifest.json") for relative_path in relative_paths: candidate = (output_dir / relative_path).resolve() if path_within(candidate, output_dir.resolve()) and candidate.exists() and candidate.is_file(): candidate.unlink() def build_export_preview_units( connection: sqlite3.Connection, paths: dict[str, Path], selected_rows: list[sqlite3.Row], ) -> list[dict[str, object]]: selected_rows_by_conversation: dict[int, list[sqlite3.Row]] = defaultdict(list) standalone_rows: list[sqlite3.Row] = [] selection_order_by_document_id: dict[int, int] = {} for index, row in enumerate(selected_rows): document_id = int(row["id"]) selection_order_by_document_id[document_id] = index child_kind = normalize_whitespace(str(row["child_document_kind"] or "")).lower() if row["conversation_id"] is None or child_kind == CHILD_DOCUMENT_KIND_ATTACHMENT: standalone_rows.append(row) else: selected_rows_by_conversation[int(row["conversation_id"])].append(row) units: list[dict[str, object]] = [] for row in standalone_rows: document_id = int(row["id"]) documents = load_preview_documents( connection, paths, document_ids=[document_id], include_attachment_children=True, require_dataset_membership=True, ) if not documents: continue title = conversation_preview_document_heading(documents[0]) units.append( { "unit_key": f"document:{document_id}", "unit_kind": "standalone", "title": title, "summary": "Standalone exported document preview", "documents": documents, "selected_document_ids": [document_id], "order_key": selection_order_by_document_id[document_id], "conversation_id": None, "conversation_type": None, } ) for conversation_id, conversation_rows in selected_rows_by_conversation.items(): conversation_row = connection.execute( "SELECT * FROM conversations WHERE id = ?", (conversation_id,), ).fetchone() if conversation_row is None: for row in conversation_rows: document_id = int(row["id"]) documents = load_preview_documents( connection, paths, document_ids=[document_id], include_attachment_children=True, require_dataset_membership=True, ) if not documents: continue title = conversation_preview_document_heading(documents[0]) units.append( { "unit_key": f"document:{document_id}", "unit_kind": "standalone", "title": title, "summary": "Standalone exported document preview", "documents": documents, "selected_document_ids": [document_id], "order_key": selection_order_by_document_id[document_id], "conversation_id": None, "conversation_type": None, } ) continue documents = load_preview_documents( connection, paths, conversation_id=conversation_id, require_dataset_membership=True, ) if not documents: continue conversation_type = normalize_whitespace(str(conversation_row["conversation_type"] or "")).lower() conversation_title = normalize_whitespace(str(conversation_row["display_name"] or "")) or f"Conversation {conversation_id}" if conversation_type == "email": selected_document_ids = [ int(row["id"]) for row in sorted(conversation_rows, key=lambda item: selection_order_by_document_id[int(item["id"])]) ] units.append( { "unit_key": f"conversation:{conversation_id}", "unit_kind": "email_conversation", "title": conversation_title, "summary": f"Full email conversation export ({len(documents)} messages)", "documents": documents, "selected_document_ids": selected_document_ids, "order_key": min(selection_order_by_document_id[int(document_id)] for document_id in selected_document_ids), "conversation_id": conversation_id, "conversation_type": conversation_type, } ) continue position_by_document_id = { int(document["id"]): index for index, document in enumerate(documents) } selected_positions = sorted( position_by_document_id[int(row["id"])] for row in conversation_rows if int(row["id"]) in position_by_document_id ) if not selected_positions: continue run_start = selected_positions[0] run_end = run_start for position in selected_positions[1:]: if position == run_end + 1: run_end = position continue run_documents = documents[run_start : run_end + 1] run_selected_document_ids = [ int(document["id"]) for document in run_documents if int(document["id"]) in selection_order_by_document_id ] units.append( { "unit_key": f"conversation:{conversation_id}:run:{run_start}:{run_end}", "unit_kind": "conversation_run", "title": conversation_title, "summary": f"Conversation run export ({len(run_documents)} documents)", "documents": run_documents, "selected_document_ids": run_selected_document_ids, "order_key": min(selection_order_by_document_id[int(document_id)] for document_id in run_selected_document_ids), "conversation_id": conversation_id, "conversation_type": conversation_type, "run_start_index": run_start, "run_end_index": run_end, } ) run_start = position run_end = position run_documents = documents[run_start : run_end + 1] run_selected_document_ids = [ int(document["id"]) for document in run_documents if int(document["id"]) in selection_order_by_document_id ] units.append( { "unit_key": f"conversation:{conversation_id}:run:{run_start}:{run_end}", "unit_kind": "conversation_run", "title": conversation_title, "summary": f"Conversation run export ({len(run_documents)} documents)", "documents": run_documents, "selected_document_ids": run_selected_document_ids, "order_key": min(selection_order_by_document_id[int(document_id)] for document_id in run_selected_document_ids), "conversation_id": conversation_id, "conversation_type": conversation_type, "run_start_index": run_start, "run_end_index": run_end, } ) units.sort(key=lambda unit: (int(unit["order_key"]), str(unit["unit_key"]))) return units def build_export_context( connection: sqlite3.Connection, rows: list[sqlite3.Row], field_defs: list[dict[str, str]], ) -> dict[str, object]: requested_field_names = {field_def["field_name"] for field_def in field_defs} context: dict[str, object] = { "dataset_memberships": {}, "production_names": {}, "attachment_parent_ids": set(), } if "dataset_id" in requested_field_names or "dataset_name" in requested_field_names: context["dataset_memberships"] = fetch_document_dataset_memberships(connection, rows) if "production_name" in requested_field_names: context["production_names"] = fetch_production_names(connection, rows) if "has_attachments" in requested_field_names: context["attachment_parent_ids"] = fetch_attachment_parent_ids( connection, [int(row["id"]) for row in rows if row["parent_document_id"] is None], ) return context def export_field_value( row: sqlite3.Row, field_def: dict[str, str], context: dict[str, object], ) -> object: field_name = field_def["field_name"] document_id = int(row["id"]) if field_name == "custodian": return document_custodian_display_text_from_row(row) if field_name == "dataset_name": dataset_membership = context["dataset_memberships"].get(document_id, {"names": []}) return "; ".join(str(dataset_name) for dataset_name in dataset_membership["names"]) if field_name == "dataset_id": dataset_membership = context["dataset_memberships"].get(document_id, {"ids": []}) membership_ids = [str(int(dataset_id)) for dataset_id in dataset_membership["ids"]] if membership_ids: return "; ".join(membership_ids) return row["dataset_id"] if field_name == "production_name": if row["production_id"] is None: return None return context["production_names"].get(int(row["production_id"])) if field_name == "is_attachment": return is_attachment_row(row) if field_name == "has_attachments": return int(row["id"]) in context["attachment_parent_ids"] return row[field_name] def serialize_export_cell_value(value: object, field_type: str) -> str: if value is None: return "" if isinstance(value, (list, tuple)): return "; ".join(str(item) for item in value if item is not None) if field_type == "boolean": return "true" if bool(value) else "false" return str(value) def resolve_export_csv_selection( connection: sqlite3.Connection, paths: dict[str, Path], document_ids: list[int] | None, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, select_from_scope: bool = False, ) -> tuple[list[sqlite3.Row], dict[str, object]]: normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) if normalized_document_ids and (query.strip() or raw_filters or sort_field or order or select_from_scope): raise RetrieverError("export-csv accepts either --doc-id selectors or query/filter/scope selectors, not both.") if normalized_document_ids: rows = fetch_visible_document_rows_by_ids(connection, normalized_document_ids) selector: dict[str, object] = { "mode": "document_ids", "document_ids": normalized_document_ids, } else: if select_from_scope: session_state = read_session_state(paths) merged_scope = merge_scope_with_search_inputs(session_state.get("scope"), query, raw_filters) selection = resolve_scope_document_search_with_explicit_sort(connection, merged_scope, sort_field, order) selector = { "mode": "scope_search", "selected_from_scope": True, "scope": selection["scope"], "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], } else: selection = resolve_document_search(connection, query, raw_filters, sort_field, order) selector = { "mode": "search", "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], } rows = [item["row"] for item in selection["results"]] return rows, selector def resolve_export_entity_table_rows( root: Path, *, query: str | None, sort_specs: list[tuple[str, str]] | None, include_ignored: bool, ) -> tuple[list[dict[str, object]], dict[str, object]]: rows: list[dict[str, object]] = [] offset = 0 limit = 200 last_payload: dict[str, object] = {} while True: payload = list_entities( root, query=query, limit=limit, offset=offset, sort_specs=sort_specs, include_ignored=include_ignored, ) last_payload = payload page_rows = [dict(item) for item in payload.get("entities", []) if isinstance(item, dict)] for item in page_rows: item["entity_status"] = item.get("canonical_status") rows.extend(page_rows) total_hits = int(payload.get("total_hits") or payload.get("total") or 0) if len(rows) >= total_hits or not page_rows: break offset += limit selector = { "mode": "table", "table": "entities", "query": last_payload.get("query") or normalize_whitespace(str(query or "")), "sort": last_payload.get("sort"), "order": last_payload.get("order"), "sort_spec": last_payload.get("sort_spec"), "include_ignored": bool(include_ignored), } return rows, selector def resolve_export_conversation_table_rows( connection: sqlite3.Connection, paths: dict[str, Path], raw_scope: object, *, sort_specs: list[tuple[str, str]] | None, ) -> tuple[list[dict[str, object]], dict[str, object]]: selection = resolve_paged_scope_conversation_search( connection, paths, raw_scope, sort_specs=sort_specs, offset=0, per_page=1_000_000_000, ) rows = [dict(item) for item in selection["results"] if isinstance(item, dict)] selector = { "mode": "table", "table": "conversations", "scope": selection["scope"], "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], "sort_spec": selection["sort_spec"], } return rows, selector def resolve_export_table_csv_selection( root: Path, connection: sqlite3.Connection, paths: dict[str, Path], session_state: dict[str, object], table_name: str, document_ids: list[int] | None, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, select_from_scope: bool = False, include_ignored_entities: bool = False, limit: int | None = None, ) -> tuple[list[object], dict[str, object]]: normalized_table_name = normalize_export_table_name(table_name) normalized_limit = None if limit is None else max(1, int(limit)) if normalized_table_name == "documents": rows, selector = resolve_export_csv_selection( connection, paths, document_ids, query, raw_filters, sort_field, order, select_from_scope, ) if normalized_limit is not None: rows = rows[:normalized_limit] selector["limit"] = normalized_limit selector["table"] = "documents" return rows, selector if document_ids: raise RetrieverError(f"/export table {normalized_table_name} does not accept --doc-id.") if normalized_table_name == "entities": if raw_filters: raise RetrieverError("/export table entities does not support --filter yet. Use an entity query instead.") effective_query = normalize_whitespace(str(query or "")) include_ignored = bool(include_ignored_entities) if select_from_scope and not effective_query: scoped_query, scoped_include_ignored = entity_browsing_options(session_state) effective_query = scoped_query or "" include_ignored = include_ignored or scoped_include_ignored if sort_field or order: sort_specs = normalize_entity_list_sort_specs(sort=sort_field or "document_count", order=order) elif select_from_scope: sort_specs = session_sort_specs(session_state, browse_mode=BROWSE_MODE_ENTITIES) or None else: sort_specs = None rows, selector = resolve_export_entity_table_rows( root, query=effective_query or None, sort_specs=sort_specs, include_ignored=include_ignored, ) selector["selected_from_scope"] = bool(select_from_scope) if normalized_limit is not None: rows = rows[:normalized_limit] selector["limit"] = normalized_limit return rows, selector effective_scope = build_effective_scope_selector( connection, paths, query=query, raw_bates=None, raw_filters=raw_filters, dataset_names=None, from_run_id=None, select_from_scope=select_from_scope, ) if sort_field or order: sort_specs = normalize_conversation_list_sort_specs(connection, sort=sort_field, order=order) elif select_from_scope: sort_specs = session_sort_specs(session_state, browse_mode=BROWSE_MODE_CONVERSATIONS) or None else: sort_specs = None rows, selector = resolve_export_conversation_table_rows( connection, paths, effective_scope, sort_specs=sort_specs, ) selector["selected_from_scope"] = bool(select_from_scope) if normalized_limit is not None: rows = rows[:normalized_limit] selector["limit"] = normalized_limit return rows, selector def normalize_export_table_limit(limit: int | None) -> int | None: return None if limit is None else max(1, int(limit)) def build_export_table_csv_seed_selector( connection: sqlite3.Connection, paths: dict[str, Path], session_state: dict[str, object], table_name: str, *, document_ids: list[int] | None, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, select_from_scope: bool, include_ignored_entities: bool, limit: int | None, ) -> dict[str, object]: normalized_table_name = normalize_export_table_name(table_name) normalized_limit = normalize_export_table_limit(limit) if document_ids: raise RetrieverError(f"/export table {normalized_table_name} does not accept --doc-id.") if normalized_table_name == "entities": if raw_filters: raise RetrieverError("/export table entities does not support --filter yet. Use an entity query instead.") effective_query = normalize_whitespace(str(query or "")) include_ignored = bool(include_ignored_entities) if select_from_scope and not effective_query: scoped_query, scoped_include_ignored = entity_browsing_options(session_state) effective_query = scoped_query or "" include_ignored = include_ignored or scoped_include_ignored if sort_field or order: sort_specs = normalize_entity_list_sort_specs(sort=sort_field or "document_count", order=order) elif select_from_scope: sort_specs = session_sort_specs(session_state, browse_mode=BROWSE_MODE_ENTITIES) or None else: sort_specs = None normalized_sort_specs = normalize_entity_list_sort_specs(sort_specs=sort_specs) if sort_specs else normalize_entity_list_sort_specs() selector: dict[str, object] = { "mode": "table", "table": "entities", "query": effective_query, "sort": normalized_sort_specs[0][0], "order": normalized_sort_specs[0][1], "sort_spec": entity_list_sort_spec_text(normalized_sort_specs), "sort_specs": serialize_sort_specs(normalized_sort_specs), "include_ignored": include_ignored, "selected_from_scope": bool(select_from_scope), } if normalized_limit is not None: selector["limit"] = normalized_limit return selector effective_scope = build_effective_scope_selector( connection, paths, query=query, raw_bates=None, raw_filters=raw_filters, dataset_names=None, from_run_id=None, select_from_scope=select_from_scope, ) if sort_field or order: sort_specs = normalize_conversation_list_sort_specs(connection, sort=sort_field, order=order) elif select_from_scope: sort_specs = session_sort_specs(session_state, browse_mode=BROWSE_MODE_CONVERSATIONS) or None else: sort_specs = None selector = { "mode": "table", "table": "conversations", "scope": effective_scope, "sort_specs": serialize_sort_specs(sort_specs), "selected_from_scope": bool(select_from_scope), } if normalized_limit is not None: selector["limit"] = normalized_limit return selector def export_csv( root: Path, raw_output_path: str, raw_fields: list[str] | None, document_ids: list[int] | None, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, select_from_scope: bool = False, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) if progress_callback is not None: progress_callback( { "status": "running", "phase": "preparing", "detail": "selecting documents", }, True, ) field_defs = resolve_export_field_definitions(connection, raw_fields) output_path = resolve_export_output_path(paths, raw_output_path) rows, selector = resolve_export_csv_selection( connection, paths, document_ids, query, raw_filters, sort_field, order, select_from_scope, ) output_path.parent.mkdir(parents=True, exist_ok=True) overwrote_existing_file = output_path.exists() context = build_export_context(connection, rows, field_defs) total_rows = len(rows) if progress_callback is not None: progress_callback( { "status": "running" if total_rows else "completed", "phase": "exporting" if total_rows else "completed", "detail": f"0/{total_rows} rows", }, True, ) with output_path.open("w", encoding="utf-8", newline="") as handle: writer = csv.writer(handle) writer.writerow([field_def["field_name"] for field_def in field_defs]) for row_index, row in enumerate(rows, start=1): writer.writerow( [ serialize_export_cell_value( export_field_value(row, field_def, context), field_def["field_type"], ) for field_def in field_defs ] ) if progress_callback is not None: progress_callback( { "status": "running", "phase": "exporting", "detail": f"{row_index}/{total_rows} rows", }, False, ) output_rel_path = None try: output_rel_path = output_path.relative_to(root).as_posix() except ValueError: output_rel_path = None payload = { "status": "ok", "output_path": str(output_path), "output_rel_path": output_rel_path, "document_count": total_rows, "field_count": len(field_defs), "fields": field_defs, "selector": selector, "overwrote_existing_file": overwrote_existing_file, "file_size": file_size_bytes(output_path), } if progress_callback is not None: progress_callback( { "status": "completed", "phase": "completed", "detail": f"{total_rows}/{total_rows} rows", }, True, ) return payload finally: connection.close() def normalize_archive_member_path(raw_rel_path: str, *, label: str = "Archive member path") -> str: normalized = normalize_whitespace(raw_rel_path) if not normalized: raise RetrieverError(f"{label} cannot be empty.") candidate = Path(normalized) if candidate.is_absolute() or ".." in candidate.parts: raise RetrieverError(f"{label} must stay within the archive root: {raw_rel_path!r}") return candidate.as_posix() def add_archive_file_once( archive: zipfile.ZipFile, written_member_paths: set[str], source_path: Path, archive_rel_path: str, ) -> str: member_path = normalize_archive_member_path(archive_rel_path) if member_path in written_member_paths: return member_path archive.write(source_path, arcname=member_path) written_member_paths.add(member_path) return member_path def add_archive_bytes_once( archive: zipfile.ZipFile, written_member_paths: set[str], payload_bytes: bytes, archive_rel_path: str, ) -> str: member_path = normalize_archive_member_path(archive_rel_path) if member_path in written_member_paths: return member_path archive.writestr(member_path, payload_bytes) written_member_paths.add(member_path) return member_path def document_preview_rows(connection: sqlite3.Connection, document_id: int) -> list[sqlite3.Row]: return connection.execute( """ SELECT rel_preview_path, preview_type, label, ordinal, created_at FROM document_previews WHERE document_id = ? ORDER BY ordinal ASC, id ASC """, (document_id,), ).fetchall() def document_source_part_rows(connection: sqlite3.Connection, document_id: int) -> list[sqlite3.Row]: return connection.execute( """ SELECT part_kind, rel_source_path, ordinal, label, created_at FROM document_source_parts WHERE document_id = ? ORDER BY CASE part_kind WHEN 'native' THEN 0 WHEN 'image' THEN 1 ELSE 2 END ASC, ordinal ASC, id ASC """, (document_id,), ).fetchall() def document_text_revision_rows(connection: sqlite3.Connection, document_id: int) -> list[sqlite3.Row]: return connection.execute( """ SELECT * FROM text_revisions WHERE document_id = ? ORDER BY id ASC """, (document_id,), ).fetchall() def document_source_text_body( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, ) -> str | None: revision_ids: list[int] = [] for field_name in ("source_text_revision_id", "active_search_text_revision_id"): if field_name in row.keys() and row[field_name] is not None: revision_id = int(row[field_name]) if revision_id not in revision_ids: revision_ids.append(revision_id) for revision_id in revision_ids: revision_row = connection.execute( """ SELECT storage_rel_path FROM text_revisions WHERE id = ? """, (revision_id,), ).fetchone() if revision_row is None: continue revision_body = read_text_revision_body(paths, revision_row["storage_rel_path"]) if revision_body is not None: return revision_body chunk_rows = document_chunk_rows(connection, int(row["id"])) if not chunk_rows: return None return "\n\n".join(str(chunk_row["text_content"] or "") for chunk_row in chunk_rows) def build_synthetic_document_export_payload( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, *, preview_rows: list[sqlite3.Row], source_part_rows: list[sqlite3.Row], ) -> bytes: payload = { "document_id": int(row["id"]), "control_number": row["control_number"], "rel_path": row["rel_path"], "file_name": row["file_name"], "file_type": row["file_type"], "source_kind": row["source_kind"], "source_rel_path": row["source_rel_path"], "source_item_id": row["source_item_id"], "production_id": row["production_id"], "parent_document_id": row["parent_document_id"], "custodian": document_custodian_display_text_from_row(row), "custodians": document_custodian_values_from_row(row), "metadata": { "author": row["author"], "content_type": row["content_type"], "custodian": document_custodian_display_text_from_row(row), "custodians": document_custodian_values_from_row(row), "date_created": row["date_created"], "date_modified": row["date_modified"], "participants": row["participants"], "recipients": row["recipients"], "subject": row["subject"], "title": row["title"], "updated_at": row["updated_at"], }, "preview_rel_paths": [ normalize_archive_member_path( str(preview_row["rel_preview_path"]), label=f"Preview path for document {int(row['id'])}", ) for preview_row in preview_rows ], "source_part_rel_paths": [ str(source_part_row["rel_source_path"]) for source_part_row in source_part_rows ], "text_content": document_source_text_body(connection, paths, row), } return (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode("utf-8") def archive_document_files( archive: zipfile.ZipFile, written_member_paths: set[str], connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, ) -> tuple[dict[str, object], list[str]]: document_id = int(row["id"]) rel_path = normalize_archive_member_path(str(row["rel_path"]), label=f"Document {document_id} rel_path") preview_rows = document_preview_rows(connection, document_id) source_part_rows = document_source_part_rows(connection, document_id) warnings: list[str] = [] exported_rel_paths: list[str] = [] preview_rel_paths: list[str] = [] source_part_entries: list[dict[str, object]] = [] source_path = resolve_workspace_artifact_path(paths["root"], str(row["rel_path"])) if source_path is not None and source_path.exists(): exported_rel_paths.append(add_archive_file_once(archive, written_member_paths, source_path, rel_path)) document_entry_kind = "copied" else: descriptor_bytes = build_synthetic_document_export_payload( connection, paths, row, preview_rows=preview_rows, source_part_rows=source_part_rows, ) exported_rel_paths.append(add_archive_bytes_once(archive, written_member_paths, descriptor_bytes, rel_path)) document_entry_kind = "synthetic" if str(row["source_kind"] or "") not in {PRODUCTION_SOURCE_KIND, PST_SOURCE_KIND, MBOX_SOURCE_KIND}: warnings.append( f"Document {document_id} has no on-disk primary artifact at {row['rel_path']}; wrote a synthetic descriptor instead." ) for preview_row in preview_rows: preview_source_path = paths["state_dir"] / str(preview_row["rel_preview_path"]) archive_rel_path = normalize_archive_member_path( str(preview_row["rel_preview_path"]), label=f"Preview path for document {document_id}", ) if not preview_source_path.exists(): warnings.append( f"Preview artifact is missing for document {document_id}: {archive_rel_path}" ) continue preview_rel_paths.append(add_archive_file_once(archive, written_member_paths, preview_source_path, archive_rel_path)) if archive_rel_path not in exported_rel_paths: exported_rel_paths.append(archive_rel_path) for source_part_row in source_part_rows: source_rel_path = normalize_archive_member_path( str(source_part_row["rel_source_path"]), label=f"Source part path for document {document_id}", ) source_part_entry = { "part_kind": str(source_part_row["part_kind"]), "rel_path": source_rel_path, "ordinal": int(source_part_row["ordinal"] or 0), "label": source_part_row["label"], } source_part_path = resolve_workspace_artifact_path(paths["root"], str(source_part_row["rel_source_path"])) if source_part_path is None or not source_part_path.exists(): source_part_entry["missing"] = True warnings.append( f"Source part is missing for document {document_id}: {source_rel_path}" ) source_part_entries.append(source_part_entry) continue source_part_entry["archive_rel_path"] = add_archive_file_once( archive, written_member_paths, source_part_path, source_rel_path, ) if source_part_entry["archive_rel_path"] not in exported_rel_paths: exported_rel_paths.append(str(source_part_entry["archive_rel_path"])) source_part_entries.append(source_part_entry) return ( { "document_id": document_id, "control_number": row["control_number"], "file_name": row["file_name"], "file_type": row["file_type"], "rel_path": rel_path, "source_kind": row["source_kind"], "parent_document_id": row["parent_document_id"], "document_entry_kind": document_entry_kind, "preview_rel_paths": preview_rel_paths, "source_part_entries": source_part_entries, "exported_rel_paths": exported_rel_paths, }, warnings, ) def archive_document_text_revisions( archive: zipfile.ZipFile, written_member_paths: set[str], connection: sqlite3.Connection, paths: dict[str, Path], document_id: int, ) -> tuple[list[dict[str, object]], list[str]]: entries: list[dict[str, object]] = [] warnings: list[str] = [] for revision_row in document_text_revision_rows(connection, document_id): entry = { "id": int(revision_row["id"]), "revision_kind": revision_row["revision_kind"], "storage_rel_path": revision_row["storage_rel_path"], } storage_rel_path = normalize_whitespace(str(revision_row["storage_rel_path"] or "")) if not storage_rel_path: entries.append(entry) continue source_path = paths["state_dir"] / storage_rel_path archive_rel_path = normalize_archive_member_path( str(Path(".retriever") / storage_rel_path), label=f"Text revision path for document {document_id}", ) if not source_path.exists(): warnings.append( f"Text revision body is missing for document {document_id}: {archive_rel_path}" ) entries.append(entry) continue entry["archive_rel_path"] = add_archive_file_once( archive, written_member_paths, source_path, archive_rel_path, ) entries.append(entry) return entries, warnings def sqlite_row_to_dict(row: sqlite3.Row) -> dict[str, object]: return {str(key): row[key] for key in row.keys()} def table_column_names_in_order(connection: sqlite3.Connection, table_name: str) -> list[str]: return [str(row["name"]) for row in table_info(connection, table_name)] def insert_row_dicts( connection: sqlite3.Connection, table_name: str, row_dicts: list[dict[str, object]], *, column_names: list[str] | None = None, ) -> None: if not row_dicts: return ordered_columns = column_names or table_column_names_in_order(connection, table_name) usable_columns = [column_name for column_name in ordered_columns if any(column_name in row_dict for row_dict in row_dicts)] if not usable_columns: return placeholders = ", ".join("?" for _ in usable_columns) connection.executemany( f""" INSERT INTO {quote_identifier(table_name)} ({', '.join(quote_identifier(column_name) for column_name in usable_columns)}) VALUES ({placeholders}) """, [[row_dict.get(column_name) for column_name in usable_columns] for row_dict in row_dicts], ) def portable_context_parent_rows( connection: sqlite3.Connection, selected_document_rows: list[sqlite3.Row], ) -> list[sqlite3.Row]: selected_document_ids = {int(row["id"]) for row in selected_document_rows} parent_ids = sorted( { int(row["parent_document_id"]) for row in selected_document_rows if row["parent_document_id"] is not None and int(row["parent_document_id"]) not in selected_document_ids } ) if not parent_ids: return [] placeholders = ", ".join("?" for _ in parent_ids) return connection.execute( f""" SELECT * FROM documents WHERE id IN ({placeholders}) ORDER BY id ASC """, parent_ids, ).fetchall() def portable_document_row_dict( row: sqlite3.Row, *, preserve_text_revisions: bool, ) -> dict[str, object]: payload = sqlite_row_to_dict(row) if not preserve_text_revisions: for field_name in ( "source_text_revision_id", "active_search_text_revision_id", "active_text_source_kind", "active_text_language", "active_text_quality_score", ): if field_name in payload: payload[field_name] = None return payload def build_portable_workspace_db( source_connection: sqlite3.Connection, portable_root: Path, selected_document_rows: list[sqlite3.Row], ) -> dict[str, object]: portable_paths = workspace_paths(portable_root) ensure_layout(portable_paths) target_connection = sqlite3.connect(portable_paths["db_path"]) target_connection.row_factory = sqlite3.Row target_connection.execute("PRAGMA busy_timeout = 5000") target_connection.execute("PRAGMA foreign_keys = ON") try: apply_schema(target_connection, portable_root) custom_field_rows = source_connection.execute( """ SELECT * FROM custom_fields_registry ORDER BY id ASC """ ).fetchall() for custom_field_row in custom_field_rows: field_name = str(custom_field_row["field_name"]) sql_type = REGISTRY_FIELD_TYPES.get(str(custom_field_row["field_type"])) if sql_type is None: raise RetrieverError(f"Unsupported custom field type for portable export: {custom_field_row['field_type']!r}") target_connection.execute( f"ALTER TABLE documents ADD COLUMN {quote_identifier(field_name)} {sql_type}" ) insert_row_dicts( target_connection, "custom_fields_registry", [sqlite_row_to_dict(row) for row in custom_field_rows], ) selected_document_ids = sorted({int(row["id"]) for row in selected_document_rows}) context_parent_rows = portable_context_parent_rows(source_connection, selected_document_rows) retained_row_by_id = { int(row["id"]): row for row in [*context_parent_rows, *selected_document_rows] } retained_rows = sorted( retained_row_by_id.values(), key=lambda row: (0 if row["parent_document_id"] is None else 1, int(row["id"])), ) stub_document_ids = sorted( document_id for document_id in retained_row_by_id if document_id not in selected_document_ids ) conversation_ids = sorted( { int(row["conversation_id"]) for row in retained_rows if row["conversation_id"] is not None } ) dataset_document_rows = source_connection.execute( f""" SELECT * FROM dataset_documents WHERE document_id IN ({', '.join('?' for _ in selected_document_ids)}) ORDER BY id ASC """ if selected_document_ids else """ SELECT * FROM dataset_documents WHERE 0 """, selected_document_ids, ).fetchall() dataset_ids = sorted( { int(row["dataset_id"]) for row in retained_rows if row["dataset_id"] is not None } | { int(row["dataset_id"]) for row in dataset_document_rows if row["dataset_id"] is not None } ) dataset_rows = source_connection.execute( f""" SELECT * FROM datasets WHERE id IN ({', '.join('?' for _ in dataset_ids)}) ORDER BY id ASC """ if dataset_ids else """ SELECT * FROM datasets WHERE 0 """, dataset_ids, ).fetchall() dataset_source_rows = source_connection.execute( f""" SELECT * FROM dataset_sources WHERE dataset_id IN ({', '.join('?' for _ in dataset_ids)}) ORDER BY id ASC """ if dataset_ids else """ SELECT * FROM dataset_sources WHERE 0 """, dataset_ids, ).fetchall() conversation_rows = source_connection.execute( f""" SELECT * FROM conversations WHERE id IN ({', '.join('?' for _ in conversation_ids)}) ORDER BY id ASC """ if conversation_ids else """ SELECT * FROM conversations WHERE 0 """, conversation_ids, ).fetchall() production_ids = sorted({int(row["production_id"]) for row in retained_rows if row["production_id"] is not None}) production_rows = source_connection.execute( f""" SELECT * FROM productions WHERE id IN ({', '.join('?' for _ in production_ids)}) ORDER BY id ASC """ if production_ids else """ SELECT * FROM productions WHERE 0 """, production_ids, ).fetchall() container_pairs = sorted( { (str(row["source_kind"]), str(row["source_rel_path"])) for row in retained_rows if normalize_whitespace(str(row["source_kind"] or "")).lower() in {PST_SOURCE_KIND, MBOX_SOURCE_KIND} and normalize_whitespace(str(row["source_rel_path"] or "")) } ) container_source_rows: list[sqlite3.Row] = [] for source_kind, source_rel_path in container_pairs: row = source_connection.execute( """ SELECT * FROM container_sources WHERE source_kind = ? AND source_rel_path = ? """, (source_kind, source_rel_path), ).fetchone() if row is not None: container_source_rows.append(row) preview_rows = source_connection.execute( f""" SELECT * FROM document_previews WHERE document_id IN ({', '.join('?' for _ in selected_document_ids)}) ORDER BY id ASC """ if selected_document_ids else """ SELECT * FROM document_previews WHERE 0 """, selected_document_ids, ).fetchall() source_part_rows = source_connection.execute( f""" SELECT * FROM document_source_parts WHERE document_id IN ({', '.join('?' for _ in selected_document_ids)}) ORDER BY id ASC """ if selected_document_ids else """ SELECT * FROM document_source_parts WHERE 0 """, selected_document_ids, ).fetchall() chunk_rows = source_connection.execute( f""" SELECT * FROM document_chunks WHERE document_id IN ({', '.join('?' for _ in selected_document_ids)}) ORDER BY id ASC """ if selected_document_ids else """ SELECT * FROM document_chunks WHERE 0 """, selected_document_ids, ).fetchall() text_revision_rows = source_connection.execute( f""" SELECT * FROM text_revisions WHERE document_id IN ({', '.join('?' for _ in selected_document_ids)}) ORDER BY id ASC """ if selected_document_ids else """ SELECT * FROM text_revisions WHERE 0 """, selected_document_ids, ).fetchall() text_revision_ids = [int(row["id"]) for row in text_revision_rows] text_revision_segment_rows = source_connection.execute( f""" SELECT * FROM text_revision_segments WHERE revision_id IN ({', '.join('?' for _ in text_revision_ids)}) ORDER BY id ASC """ if text_revision_ids else """ SELECT * FROM text_revision_segments WHERE 0 """, text_revision_ids, ).fetchall() target_connection.execute("PRAGMA foreign_keys = OFF") target_connection.execute("BEGIN") try: insert_row_dicts(target_connection, "conversations", [sqlite_row_to_dict(row) for row in conversation_rows]) insert_row_dicts(target_connection, "datasets", [sqlite_row_to_dict(row) for row in dataset_rows]) insert_row_dicts(target_connection, "dataset_sources", [sqlite_row_to_dict(row) for row in dataset_source_rows]) insert_row_dicts(target_connection, "productions", [sqlite_row_to_dict(row) for row in production_rows]) insert_row_dicts(target_connection, "container_sources", [sqlite_row_to_dict(row) for row in container_source_rows]) insert_row_dicts( target_connection, "documents", [ portable_document_row_dict( row, preserve_text_revisions=int(row["id"]) not in set(stub_document_ids), ) for row in retained_rows ], column_names=table_column_names_in_order(target_connection, "documents"), ) insert_row_dicts(target_connection, "dataset_documents", [sqlite_row_to_dict(row) for row in dataset_document_rows]) insert_row_dicts(target_connection, "document_previews", [sqlite_row_to_dict(row) for row in preview_rows]) insert_row_dicts(target_connection, "document_source_parts", [sqlite_row_to_dict(row) for row in source_part_rows]) insert_row_dicts( target_connection, "text_revisions", [ { **sqlite_row_to_dict(row), "created_by_job_version_id": None, } for row in text_revision_rows ], ) insert_row_dicts( target_connection, "text_revision_segments", [sqlite_row_to_dict(row) for row in text_revision_segment_rows], ) insert_row_dicts(target_connection, "document_chunks", [sqlite_row_to_dict(row) for row in chunk_rows]) for document_id in selected_document_ids: refresh_documents_fts_row(target_connection, document_id) if chunk_rows: target_connection.executemany( """ INSERT INTO chunks_fts (chunk_id, document_id, text_content) VALUES (?, ?, ?) """, [ (row["id"], row["document_id"], row["text_content"]) for row in chunk_rows ], ) target_connection.commit() except Exception: target_connection.rollback() raise target_connection.execute("PRAGMA foreign_keys = ON") foreign_key_issues = target_connection.execute("PRAGMA foreign_key_check").fetchall() if foreign_key_issues: raise RetrieverError(f"Portable workspace export failed foreign key validation: {foreign_key_issues[0]}") return { "db_path": portable_paths["db_path"], "selected_document_ids": selected_document_ids, "retained_document_ids": [int(row["id"]) for row in retained_rows], "stub_document_ids": stub_document_ids, } finally: target_connection.close() def resolve_export_archive_selection( connection: sqlite3.Connection, paths: dict[str, Path], *, dataset_names: list[str] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, family_mode: str = "exact", seed_limit: int | None = None, ) -> tuple[list[dict[str, object]], dict[str, object], str]: normalized_family_mode = normalize_run_family_mode(family_mode) if seed_limit is not None and seed_limit < 1: raise RetrieverError("Archive limit must be >= 1.") selector = build_effective_scope_selector( connection, paths, query=query, raw_bates=raw_bates, raw_filters=raw_filters, dataset_names=dataset_names, from_run_id=from_run_id, select_from_scope=select_from_scope, ) if not scope_run_selector_has_inputs(selector): raise RetrieverError("Archive selector must include at least one inclusion input.") selected_documents, _ = plan_scope_selected_documents( connection, selector=selector, family_mode=normalized_family_mode, seed_limit=seed_limit, ) return selected_documents, selector, normalized_family_mode def export_archive( root: Path, raw_output_path: str, *, dataset_names: list[str] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, family_mode: str = "exact", seed_limit: int | None = None, portable_workspace: bool = False, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) metadata_field_defs = resolve_archive_metadata_field_definitions(connection) selected_documents, selector, normalized_family_mode = resolve_export_archive_selection( connection, paths, dataset_names=dataset_names, query=query, raw_bates=raw_bates, raw_filters=raw_filters, from_run_id=from_run_id, select_from_scope=select_from_scope, family_mode=family_mode, seed_limit=seed_limit, ) output_path = resolve_export_output_path(paths, raw_output_path) output_path.parent.mkdir(parents=True, exist_ok=True) overwrote_existing_file = output_path.exists() written_member_paths: set[str] = set() manifest_document_entries: list[dict[str, object]] = [] warnings: list[str] = [] created_at = utc_now() total_documents = len(selected_documents) if progress_callback is not None: progress_callback( { "status": "running" if total_documents or portable_workspace else "completed", "phase": "exporting" if total_documents else ("building portable workspace" if portable_workspace else "completed"), "detail": f"0/{total_documents} documents", }, True, ) with zipfile.ZipFile(output_path, "w", compression=zipfile.ZIP_DEFLATED) as archive: for document_index, selected_document in enumerate(selected_documents, start=1): document_row = selected_document["document_row"] manifest_entry, document_warnings = archive_document_files( archive, written_member_paths, connection, paths, document_row, ) manifest_entry["ordinal"] = int(selected_document["ordinal"]) manifest_entry["inclusion_reason"] = selected_document["inclusion_reason"] if portable_workspace: revision_entries, revision_warnings = archive_document_text_revisions( archive, written_member_paths, connection, paths, int(document_row["id"]), ) manifest_entry["text_revision_entries"] = revision_entries warnings.extend(revision_warnings) manifest_document_entries.append(manifest_entry) warnings.extend(document_warnings) if progress_callback is not None: progress_callback( { "status": "running", "phase": "exporting", "detail": f"{document_index}/{total_documents} documents", }, False, ) portable_workspace_payload = None if portable_workspace: if progress_callback is not None: progress_callback( { "status": "running", "phase": "building portable workspace", "detail": f"{total_documents}/{total_documents} documents", }, False, ) with tempfile.TemporaryDirectory(prefix="retriever-portable-workspace-") as tempdir: portable_root = Path(tempdir) / "workspace" portable_workspace_payload = build_portable_workspace_db( connection, portable_root, [selected_document["document_row"] for selected_document in selected_documents], ) add_archive_file_once( archive, written_member_paths, Path(portable_workspace_payload["db_path"]), ".retriever/retriever.db", ) if progress_callback is not None: progress_callback( { "status": "running", "phase": "finalizing", "detail": f"{total_documents}/{total_documents} documents", }, False, ) selected_rows = [selected_document["document_row"] for selected_document in selected_documents] sidecar_payload = write_archive_sidecars( archive, written_member_paths, connection, paths, selected_rows, manifest_document_entries, metadata_payload=build_archive_metadata_csv_bytes( connection, selected_rows, manifest_document_entries, metadata_field_defs, ), created_at=created_at, selector=selector, family_mode=normalized_family_mode, seed_limit=seed_limit, portable_workspace=portable_workspace, portable_workspace_payload=portable_workspace_payload, warnings=warnings, failed_count=0, ) checksum_payload = ensure_archive_checksums_file(output_path) output_rel_path = None try: output_rel_path = output_path.relative_to(root).as_posix() except ValueError: output_rel_path = None payload = { "status": "ok", "created_at": created_at, "output_path": str(output_path), "output_rel_path": output_rel_path, "document_count": total_documents, "selector": selector, "family_mode": normalized_family_mode, "seed_limit": seed_limit, "portable_workspace": portable_workspace, "manifest_rel_path": ARCHIVE_MANIFEST_REL_PATH, "root_manifest_rel_path": ARCHIVE_ROOT_MANIFEST_REL_PATH, "metadata_rel_path": ARCHIVE_METADATA_REL_PATH, "readme_rel_path": ARCHIVE_ROOT_README_REL_PATH, "preview_root_rel_path": ARCHIVE_PREVIEW_ROOT_REL_PATH, "previews_index_rel_path": ARCHIVE_PREVIEWS_INDEX_REL_PATH, "source_parts_rel_path": ARCHIVE_SOURCE_PARTS_REL_PATH, "extracted_text_index_rel_path": ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH, "text_root_rel_path": ARCHIVE_TEXT_ROOT_REL_PATH, "checksums_rel_path": ARCHIVE_CHECKSUMS_REL_PATH, "loadfile_rel_path": ARCHIVE_LOADFILE_REL_PATH, "archive_member_count": int(checksum_payload["archive_member_count"]), "documents": manifest_document_entries, "counts": sidecar_payload["sidecar_counts"], "warnings": warnings, "overwrote_existing_file": overwrote_existing_file, "file_size": file_size_bytes(output_path), } if progress_callback is not None: progress_callback( { "status": "completed", "phase": "completed", "detail": f"{total_documents}/{total_documents} documents", }, True, ) return payload finally: connection.close() EXPORT_RUN_ACTIVE_STATUSES = {"seeding", "exporting", "finalizing"} EXPORT_RUN_TERMINAL_STATUSES = {"completed", "failed", "canceled"} EXPORT_RUN_STEP_MIN_REMAINING_SECONDS = 1.25 EXPORT_CSV_ROW_BATCH_SIZE = 100 EXPORT_CSV_FINALIZE_FRAGMENT_BATCH_SIZE = 50 EXPORT_TABLE_CSV_SEED_PAGE_SIZE = 1000 def export_run_temp_dir(paths: dict[str, Path], run_id: str) -> Path: return paths["state_dir"] / "tmp" / "exports" / run_id def export_run_csv_rows_dir(paths: dict[str, Path], run_id: str) -> Path: return export_run_temp_dir(paths, run_id) / "csv_rows" def export_run_csv_fragment_rel_path(run_id: str, ordinal: int) -> str: return str(Path("tmp") / "exports" / run_id / "csv_rows" / f"{int(ordinal):08d}.csvfrag") def export_run_partial_csv_path(paths: dict[str, Path], run_id: str) -> Path: return export_run_temp_dir(paths, run_id) / "output.partial.csv" def export_run_partial_archive_path(paths: dict[str, Path], run_id: str) -> Path: return export_run_temp_dir(paths, run_id) / "output.partial.zip" def write_text_file_atomic(path: Path, text: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) temp_file = tempfile.NamedTemporaryFile( mode="w", encoding="utf-8", dir=str(path.parent), prefix=f".{path.name}.", suffix=".tmp", delete=False, ) temp_path = Path(temp_file.name) try: with temp_file: temp_file.write(text) temp_file.flush() os.replace(temp_path, path) except Exception: try: temp_path.unlink() except FileNotFoundError: pass raise def export_run_output_rel_path(root: Path, output_path: Path) -> str | None: try: return output_path.relative_to(root).as_posix() except ValueError: return None def export_run_current_output_path(root: Path, row: sqlite3.Row) -> Path: output_rel_path = normalize_inline_whitespace(str(row["output_rel_path"] or "")) if output_rel_path: return (root / output_rel_path).resolve() return Path(str(row["output_path"])) def latest_export_run_row( connection: sqlite3.Connection, *, export_kind: str, run_id: str | None = None, ) -> sqlite3.Row | None: if run_id: return connection.execute( "SELECT * FROM export_runs WHERE run_id = ? AND export_kind = ?", (run_id, export_kind), ).fetchone() return connection.execute( """ SELECT * FROM export_runs WHERE export_kind = ? ORDER BY id DESC LIMIT 1 """, (export_kind,), ).fetchone() def active_export_run_row(connection: sqlite3.Connection, *, export_kind: str) -> sqlite3.Row | None: return connection.execute( """ SELECT * FROM export_runs WHERE export_kind = ? AND status IN ('seeding', 'exporting', 'finalizing') ORDER BY id DESC LIMIT 1 """, (export_kind,), ).fetchone() def require_export_run_row( connection: sqlite3.Connection, *, export_kind: str, run_id: str, ) -> sqlite3.Row: row = latest_export_run_row(connection, export_kind=export_kind, run_id=run_id) if row is None: raise RetrieverError(f"Unknown {export_kind} export run: {run_id}") return row def export_work_item_counts(connection: sqlite3.Connection, *, run_id: str) -> dict[str, int]: counts = {status: 0 for status in ("pending", "running", "completed", "failed")} for row in connection.execute( """ SELECT status, COUNT(*) AS count FROM export_work_items WHERE run_id = ? GROUP BY status """, (run_id,), ): counts[str(row["status"])] = int(row["count"] or 0) return counts def export_work_item_counts_by_unit_type(connection: sqlite3.Connection, *, run_id: str) -> dict[str, dict[str, int]]: by_unit_type: dict[str, dict[str, int]] = {} for row in connection.execute( """ SELECT unit_type, status, COUNT(*) AS count FROM export_work_items WHERE run_id = ? GROUP BY unit_type, status ORDER BY unit_type ASC, status ASC """, (run_id,), ): unit_type = str(row["unit_type"]) by_unit_type.setdefault(unit_type, {status: 0 for status in ("pending", "running", "completed", "failed")}) by_unit_type[unit_type][str(row["status"])] = int(row["count"] or 0) return by_unit_type def refresh_export_run_counts(connection: sqlite3.Connection, *, run_id: str, now: str | None = None) -> None: counts = export_work_item_counts(connection, run_id=run_id) connection.execute( """ UPDATE export_runs SET completed_items = ?, failed_items = ?, last_heartbeat_at = ? WHERE run_id = ? """, (counts.get("completed", 0), counts.get("failed", 0), now or utc_now(), run_id), ) def export_run_next_commands(root: Path, row: sqlite3.Row, *, budget_seconds: int) -> list[str]: run_id = shlex.quote(str(row["run_id"])) root_arg = shlex.quote(str(root)) export_kind = str(row["export_kind"]) prefix = "export-csv" if export_kind == "csv" else "export-archive" if str(row["status"]) in EXPORT_RUN_TERMINAL_STATUSES: return [f"{prefix}-status {root_arg} --run-id {run_id}"] budget_arg = str(int(budget_seconds)) return [ f"{prefix}-run-step {root_arg} --run-id {run_id} --budget-seconds {budget_arg}", f"{prefix}-status {root_arg} --run-id {run_id}", ] def export_run_status_payload( connection: sqlite3.Connection, root: Path, row: sqlite3.Row, *, budget_seconds: int = DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, ) -> dict[str, object]: run_id = str(row["run_id"]) return { "run_id": run_id, "export_kind": row["export_kind"], "status": row["status"], "phase": row["phase"], "output_path": str(export_run_current_output_path(root, row)), "output_rel_path": row["output_rel_path"], "selector": decode_json_text(row["selector_json"], default={}) or {}, "config": decode_json_text(row["config_json"], default={}) or {}, "cursor": decode_json_text(row["cursor_json"], default={}) or {}, "total_items": int(row["total_items"] or 0), "completed_items": int(row["completed_items"] or 0), "failed_items": int(row["failed_items"] or 0), "counts": { "work_items": export_work_item_counts(connection, run_id=run_id), "by_unit_type": export_work_item_counts_by_unit_type(connection, run_id=run_id), }, "created_at": row["created_at"], "started_at": row["started_at"], "completed_at": row["completed_at"], "last_error": row["error"], "budget_recommendation_seconds": int(budget_seconds), "next_recommended_commands": export_run_next_commands(root, row, budget_seconds=budget_seconds), } def export_status( root: Path, *, export_kind: str, run_id: str | None = None, budget_seconds: int | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = latest_export_run_row(connection, export_kind=export_kind, run_id=run_id) if row is None: return { "ok": True, "status": "none", "phase": None, "run_id": None, "export_kind": export_kind, "counts": { "work_items": {status: 0 for status in ("pending", "running", "completed", "failed")}, "by_unit_type": {}, }, "next_recommended_commands": [], } return {"ok": True, **export_run_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() def export_run_facade_payload( *, root: Path, export_kind: str, budget_seconds: int, created: bool, run_payload: dict[str, object], step_payloads: list[dict[str, object]], reason: str, mode: str, ) -> dict[str, object]: more_work_remaining = str(run_payload.get("status")) not in EXPORT_RUN_TERMINAL_STATUSES executed_steps: list[str] = [] for step_payload in step_payloads: executed_steps.extend(str(step) for step in list(step_payload.get("executed_steps") or [])) next_commands: list[str] = [] if more_work_remaining: prefix = "export-csv-start" if export_kind == "csv" else "export-archive-start" run_id = run_payload.get("run_id") if run_id: next_commands.append( f"{prefix} {shlex.quote(str(root))} ... --run-to-completion --budget-seconds {int(budget_seconds)}" ) next_commands.extend(str(command) for command in list(run_payload.get("next_recommended_commands") or [])) return { "ok": True, "export_kind": export_kind, "mode": mode, "created": created, "run_id": run_payload.get("run_id"), "status": run_payload.get("status"), "phase": run_payload.get("phase"), "reason": reason, "executed": bool(executed_steps), "executed_steps": executed_steps, "step_calls": len(step_payloads), "step_results": step_payloads, "more_work_remaining": more_work_remaining, "run": run_payload, "counts": run_payload.get("counts"), "next_recommended_commands": next_commands, } def export_run_to_completion( root: Path, *, export_kind: str, run_id: str, budget_seconds: int, created: bool, start_payload: dict[str, object], progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: run_payload = { key: value for key, value in start_payload.items() if key not in {"ok", "created"} } step_payloads: list[dict[str, object]] = [] reason = "run_terminal" if str(run_payload.get("status")) in EXPORT_RUN_TERMINAL_STATUSES else "budget_exhausted" if progress_callback is not None: progress_callback(start_payload, True) while str(run_payload.get("status")) not in EXPORT_RUN_TERMINAL_STATUSES: step_payload = export_run_step( root, export_kind=export_kind, run_id=run_id, budget_seconds=budget_seconds, ) step_payloads.append(step_payload) if isinstance(step_payload.get("run"), dict): run_payload = dict(step_payload["run"]) reason = str(step_payload.get("reason") or reason) if progress_callback is not None: progress_callback(step_payload, False) if not bool(step_payload.get("executed")): break if reason == "no_runnable_step" and bool(step_payload.get("more_work_remaining")): break payload = export_run_facade_payload( root=root, export_kind=export_kind, budget_seconds=budget_seconds, created=created, run_payload=run_payload, step_payloads=step_payloads, reason=reason, mode="run_to_completion", ) if progress_callback is not None: progress_callback(payload, True) return payload def export_csv_start( root: Path, raw_output_path: str, raw_fields: list[str] | None, document_ids: list[int] | None, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, select_from_scope: bool = False, budget_seconds: int | None = None, run_to_completion: bool = False, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) start_payload: dict[str, object] | None = None run_id: str | None = None try: apply_schema(connection, root) active_row = active_export_run_row(connection, export_kind="csv") if active_row is not None: raise RetrieverError(f"CSV export run {active_row['run_id']} is already active.") field_defs = resolve_export_field_definitions(connection, raw_fields) output_path = resolve_export_output_path(paths, raw_output_path) rows, selector = resolve_export_csv_selection( connection, paths, document_ids, query, raw_filters, sort_field, order, select_from_scope, ) run_id = new_ingest_v2_run_id() now = utc_now() output_rel_path = export_run_output_rel_path(root, output_path) config = { "fields": field_defs, "raw_output_path": raw_output_path, "overwrote_existing_file": output_path.exists(), } cursor = { "assembled_ordinal": 0, "partial_rel_path": str(Path("tmp") / "exports" / run_id / "output.partial.csv"), } export_run_csv_rows_dir(paths, run_id).mkdir(parents=True, exist_ok=True) connection.execute("BEGIN") try: connection.execute( """ INSERT INTO export_runs ( run_id, export_kind, output_path, output_rel_path, selector_json, config_json, cursor_json, phase, status, total_items, completed_items, failed_items, created_at, started_at, last_heartbeat_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?) """, ( run_id, "csv", str(output_path), output_rel_path, compact_json_text(selector), compact_json_text(config), compact_json_text(cursor), "exporting", "exporting" if rows else "finalizing", len(rows), now, now, now, ), ) connection.executemany( """ INSERT INTO export_work_items ( run_id, unit_type, ordinal, document_id, payload_json, artifact_manifest_json, status, created_at, updated_at ) VALUES (?, 'csv_row', ?, ?, ?, '{}', 'pending', ?, ?) """, [ ( run_id, index, int(row["id"]), compact_json_text({"document_id": int(row["id"])}), now, now, ) for index, row in enumerate(rows, start=1) ], ) connection.commit() except Exception: connection.rollback() raise row = require_export_run_row(connection, export_kind="csv", run_id=run_id) start_payload = {"ok": True, "created": True, **export_run_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() if start_payload is None or run_id is None: raise RetrieverError("CSV export start did not produce a resumable run.") if run_to_completion: return export_run_to_completion( root, export_kind="csv", run_id=run_id, budget_seconds=budget, created=True, start_payload=start_payload, progress_callback=progress_callback, ) return start_payload def export_table_csv_start( root: Path, raw_output_path: str, *, table_name: str, raw_fields: list[str] | None, document_ids: list[int] | None = None, query: str = "", raw_filters: list[list[str]] | None = None, sort_field: str | None = None, order: str | None = None, select_from_scope: bool = False, include_ignored_entities: bool = False, limit: int | None = None, budget_seconds: int | None = None, run_to_completion: bool = False, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: normalized_table_name = normalize_export_table_name(table_name) budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) start_payload: dict[str, object] | None = None run_id: str | None = None try: apply_schema(connection, root) active_row = active_export_run_row(connection, export_kind="csv") if active_row is not None: raise RetrieverError(f"CSV export run {active_row['run_id']} is already active.") session_state = read_session_state(paths) field_defs = resolve_table_export_field_definitions(connection, normalized_table_name, raw_fields) output_path = resolve_export_output_path(paths, raw_output_path) if normalized_table_name == "documents": rows, selector = resolve_export_table_csv_selection( root, connection, paths, session_state, normalized_table_name, document_ids, query, raw_filters, sort_field, order, select_from_scope, include_ignored_entities=include_ignored_entities, limit=limit, ) initial_phase = "exporting" if rows else "finalizing" initial_total_items = len(rows) else: rows = [] selector = build_export_table_csv_seed_selector( connection, paths, session_state, normalized_table_name, document_ids=document_ids, query=query, raw_filters=raw_filters, sort_field=sort_field, order=order, select_from_scope=select_from_scope, include_ignored_entities=include_ignored_entities, limit=limit, ) initial_phase = "seeding" initial_total_items = 0 run_id = new_ingest_v2_run_id() now = utc_now() output_rel_path = export_run_output_rel_path(root, output_path) config = { "fields": field_defs, "raw_output_path": raw_output_path, "overwrote_existing_file": output_path.exists(), "table": normalized_table_name, } cursor = { "assembled_ordinal": 0, "partial_rel_path": str(Path("tmp") / "exports" / run_id / "output.partial.csv"), "seed_offset": 0, "seed_next_batch_ordinal": 1, "seed_page_size": EXPORT_TABLE_CSV_SEED_PAGE_SIZE, "seed_batch_mode": "row_batches", } export_run_csv_rows_dir(paths, run_id).mkdir(parents=True, exist_ok=True) if normalized_table_name == "documents": unit_type = "csv_row" work_items = [ ( run_id, unit_type, index, int(row["id"]), # type: ignore[index] compact_json_text({"table": "documents", "document_id": int(row["id"])}), # type: ignore[index] now, now, ) for index, row in enumerate(rows, start=1) ] else: row_id_key = "entity_id" if normalized_table_name == "entities" else "conversation_id" unit_type = "csv_entity_row" if normalized_table_name == "entities" else "csv_conversation_row" work_items = [ ( run_id, unit_type, index, None, compact_json_text( { "table": normalized_table_name, row_id_key: int(row["id"]), # type: ignore[index] "row": dict(row), # type: ignore[arg-type] } ), now, now, ) for index, row in enumerate(rows, start=1) ] connection.execute("BEGIN") try: connection.execute( """ INSERT INTO export_runs ( run_id, export_kind, output_path, output_rel_path, selector_json, config_json, cursor_json, phase, status, total_items, completed_items, failed_items, created_at, started_at, last_heartbeat_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?) """, ( run_id, "csv", str(output_path), output_rel_path, compact_json_text(selector), compact_json_text(config), compact_json_text(cursor), initial_phase, initial_phase, initial_total_items, now, now, now, ), ) connection.executemany( """ INSERT INTO export_work_items ( run_id, unit_type, ordinal, document_id, payload_json, artifact_manifest_json, status, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, '{}', 'pending', ?, ?) """, work_items, ) connection.commit() except Exception: connection.rollback() raise row = require_export_run_row(connection, export_kind="csv", run_id=run_id) start_payload = {"ok": True, "created": True, **export_run_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() if start_payload is None or run_id is None: raise RetrieverError("Table CSV export start did not produce a resumable run.") if run_to_completion: return export_run_to_completion( root, export_kind="csv", run_id=run_id, budget_seconds=budget, created=True, start_payload=start_payload, progress_callback=progress_callback, ) return start_payload def export_archive_start( root: Path, raw_output_path: str, *, dataset_names: list[str] | None = None, query: str = "", raw_bates: str | None = None, raw_filters: list[list[str]] | None = None, from_run_id: int | None = None, select_from_scope: bool = False, family_mode: str = "exact", seed_limit: int | None = None, portable_workspace: bool = False, budget_seconds: int | None = None, run_to_completion: bool = False, progress_callback: Callable[[dict[str, object], bool], None] | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) start_payload: dict[str, object] | None = None run_id: str | None = None try: apply_schema(connection, root) active_row = active_export_run_row(connection, export_kind="archive") if active_row is not None: raise RetrieverError(f"Archive export run {active_row['run_id']} is already active.") metadata_field_defs = resolve_archive_metadata_field_definitions(connection) selected_documents, selector, normalized_family_mode = resolve_export_archive_selection( connection, paths, dataset_names=dataset_names, query=query, raw_bates=raw_bates, raw_filters=raw_filters, from_run_id=from_run_id, select_from_scope=select_from_scope, family_mode=family_mode, seed_limit=seed_limit, ) output_path = resolve_export_output_path(paths, raw_output_path) run_id = new_ingest_v2_run_id() now = utc_now() output_rel_path = export_run_output_rel_path(root, output_path) config = { "raw_output_path": raw_output_path, "family_mode": normalized_family_mode, "seed_limit": seed_limit, "portable_workspace": bool(portable_workspace), "overwrote_existing_file": output_path.exists(), "manifest_rel_path": ARCHIVE_MANIFEST_REL_PATH, "root_manifest_rel_path": ARCHIVE_ROOT_MANIFEST_REL_PATH, "metadata_rel_path": ARCHIVE_METADATA_REL_PATH, "readme_rel_path": ARCHIVE_ROOT_README_REL_PATH, "preview_root_rel_path": ARCHIVE_PREVIEW_ROOT_REL_PATH, "previews_index_rel_path": ARCHIVE_PREVIEWS_INDEX_REL_PATH, "source_parts_rel_path": ARCHIVE_SOURCE_PARTS_REL_PATH, "extracted_text_index_rel_path": ARCHIVE_EXTRACTED_TEXT_INDEX_REL_PATH, "text_root_rel_path": ARCHIVE_TEXT_ROOT_REL_PATH, "checksums_rel_path": ARCHIVE_CHECKSUMS_REL_PATH, "loadfile_rel_path": ARCHIVE_LOADFILE_REL_PATH, "metadata_fields": metadata_field_defs, } cursor = { "created_at": now, "partial_rel_path": str(Path("tmp") / "exports" / run_id / "output.partial.zip"), "portable_workspace_added": False, "manifest_added": False, "metadata_added": False, } export_run_temp_dir(paths, run_id).mkdir(parents=True, exist_ok=True) connection.execute("BEGIN") try: connection.execute( """ INSERT INTO export_runs ( run_id, export_kind, output_path, output_rel_path, selector_json, config_json, cursor_json, phase, status, total_items, completed_items, failed_items, created_at, started_at, last_heartbeat_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, 0, ?, ?, ?) """, ( run_id, "archive", str(output_path), output_rel_path, compact_json_text(selector), compact_json_text(config), compact_json_text(cursor), "exporting", "exporting" if selected_documents else "finalizing", len(selected_documents), now, now, now, ), ) connection.executemany( """ INSERT INTO export_work_items ( run_id, unit_type, ordinal, document_id, payload_json, artifact_manifest_json, status, created_at, updated_at ) VALUES (?, 'archive_document', ?, ?, ?, '{}', 'pending', ?, ?) """, [ ( run_id, int(selected_document["ordinal"]), int(selected_document["document_row"]["id"]), compact_json_text( { "document_id": int(selected_document["document_row"]["id"]), "inclusion_reason": selected_document["inclusion_reason"], } ), now, now, ) for selected_document in selected_documents ], ) connection.commit() except Exception: connection.rollback() raise row = require_export_run_row(connection, export_kind="archive", run_id=run_id) start_payload = {"ok": True, "created": True, **export_run_status_payload(connection, root, row, budget_seconds=budget)} finally: connection.close() if start_payload is None or run_id is None: raise RetrieverError("Archive export start did not produce a resumable run.") if run_to_completion: return export_run_to_completion( root, export_kind="archive", run_id=run_id, budget_seconds=budget, created=True, start_payload=start_payload, progress_callback=progress_callback, ) return start_payload def export_table_csv_seed_page( root: Path, connection: sqlite3.Connection, paths: dict[str, Path], table_name: str, selector: dict[str, object], *, offset: int, page_size: int, known_total_hits: int | None = None, ) -> tuple[list[dict[str, object]], int]: normalized_table_name = normalize_export_table_name(table_name) if normalized_table_name == "entities": sort_specs = coerce_sort_specs(selector.get("sort_specs")) or None normalized_sort_specs = normalize_entity_list_sort_specs(sort_specs=sort_specs) if sort_specs else normalize_entity_list_sort_specs() normalized_query = normalize_whitespace(str(selector.get("query") or "")).lower() include_ignored = bool(selector.get("include_ignored")) where_clauses = ["e.canonical_status != ?"] if include_ignored else ["e.canonical_status = ?"] params: list[object] = [ENTITY_STATUS_MERGED if include_ignored else ENTITY_STATUS_ACTIVE] if normalized_query: like_query = f"%{normalized_query}%" where_clauses.append( """ ( LOWER(COALESCE(e.display_name, '')) LIKE ? OR LOWER(COALESCE(e.primary_email, '')) LIKE ? OR LOWER(COALESCE(e.primary_phone, '')) LIKE ? OR EXISTS ( SELECT 1 FROM entity_identifiers ei WHERE ei.entity_id = e.id AND ( LOWER(COALESCE(ei.display_value, '')) LIKE ? OR LOWER(COALESCE(ei.normalized_value, '')) LIKE ? ) ) ) """ ) params.extend([like_query, like_query, like_query, like_query, like_query]) where_sql = " AND ".join(where_clauses) if known_total_hits is None: total_row = connection.execute( f""" SELECT COUNT(*) AS total FROM entities e WHERE {where_sql} """, tuple(params), ).fetchone() total_hits = int(total_row["total"] if total_row is not None else 0) else: total_hits = max(0, int(known_total_hits)) rows = connection.execute( f""" SELECT e.*, COUNT(DISTINCT de.document_id) AS document_count, GROUP_CONCAT(DISTINCT de.role) AS roles FROM entities e LEFT JOIN document_entities de ON de.entity_id = e.id WHERE {where_sql} GROUP BY e.id ORDER BY {entity_list_order_sql(normalized_sort_specs)} LIMIT ? OFFSET ? """, (*params, max(1, int(page_size)), max(0, int(offset))), ).fetchall() entity_ids = [int(row["id"]) for row in rows] identifiers_by_entity = entity_identifiers_by_entity_id(connection, entity_ids) entities = [ serialize_entity_summary(row, identifiers_by_entity.get(int(row["id"]), [])) for row in rows ] for item in entities: item["entity_status"] = item.get("canonical_status") return entities, total_hits sort_specs = coerce_sort_specs(selector.get("sort_specs")) or None raw_scope = coerce_scope_payload(selector.get("scope")) selection = resolve_paged_scope_conversation_search( connection, paths, raw_scope, sort_specs=sort_specs, offset=offset, per_page=page_size, ) rows = [dict(item) for item in selection.get("results", []) if isinstance(item, dict)] total_hits = int(selection.get("total_hits") or 0) return rows, total_hits def export_table_csv_seed_step( root: Path, connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, *, deadline: float, ) -> dict[str, object]: run_id = str(row["run_id"]) config = decode_json_text(row["config_json"], default={}) or {} selector = decode_json_text(row["selector_json"], default={}) or {} cursor = decode_json_text(row["cursor_json"], default={}) or {} table_name = normalize_export_table_name(str(config.get("table") or selector.get("table") or "documents")) if table_name == "documents": raise RetrieverError("Document CSV exports do not use table seeding.") if cursor.get("seed_batch_mode") != "row_batches": connection.execute("DELETE FROM export_work_items WHERE run_id = ?", (run_id,)) cursor["seed_offset"] = 0 cursor["seed_next_batch_ordinal"] = 1 cursor["seed_page_size"] = EXPORT_TABLE_CSV_SEED_PAGE_SIZE cursor["seed_batch_mode"] = "row_batches" cursor.pop("seed_next_ordinal", None) connection.execute( """ UPDATE export_runs SET total_items = 0, completed_items = 0, failed_items = 0, cursor_json = ?, last_heartbeat_at = ? WHERE run_id = ? """, (compact_json_text(cursor), utc_now(), run_id), ) connection.commit() seed_offset = int(cursor.get("seed_offset") or 0) next_batch_ordinal = int(cursor.get("seed_next_batch_ordinal") or 1) page_size = max(1, int(cursor.get("seed_page_size") or EXPORT_TABLE_CSV_SEED_PAGE_SIZE)) known_total_hits = int(cursor["seed_total_hits"]) if cursor.get("seed_total_hits") is not None else None explicit_limit = normalize_export_table_limit(int(selector["limit"])) if selector.get("limit") is not None else None seeded_rows = 0 seeded_batches = 0 total_to_export = int(row["total_items"] or 0) total_rows = int(cursor.get("seed_total_rows") or 0) completed = False while ingest_v2_deadline_remaining_seconds(deadline) >= EXPORT_RUN_STEP_MIN_REMAINING_SECONDS: if explicit_limit is not None and seed_offset >= explicit_limit: total_rows = explicit_limit completed = True break remaining_limit = None if explicit_limit is None else max(0, explicit_limit - seed_offset) if remaining_limit == 0: total_rows = explicit_limit or total_rows completed = True break current_page_size = page_size if remaining_limit is None else min(page_size, remaining_limit) page_rows, total_hits = export_table_csv_seed_page( root, connection, paths, table_name, selector, offset=seed_offset, page_size=current_page_size, known_total_hits=known_total_hits, ) known_total_hits = total_hits total_rows = min(total_hits, explicit_limit) if explicit_limit is not None else total_hits total_to_export = (total_rows + page_size - 1) // page_size if total_rows else 0 if not page_rows: completed = True break row_id_key = "entity_id" if table_name == "entities" else "conversation_id" unit_type = "csv_entity_batch" if table_name == "entities" else "csv_conversation_batch" now = utc_now() connection.execute( """ INSERT INTO export_work_items ( run_id, unit_type, ordinal, document_id, payload_json, artifact_manifest_json, status, created_at, updated_at ) VALUES (?, ?, ?, NULL, ?, '{}', 'pending', ?, ?) """, ( run_id, unit_type, next_batch_ordinal, compact_json_text( { "table": table_name, "row_id_key": row_id_key, "offset": seed_offset, "row_count": len(page_rows), "rows": [dict(page_row) for page_row in page_rows], } ), now, now, ), ) seeded_rows += len(page_rows) seeded_batches += 1 seed_offset += len(page_rows) next_batch_ordinal += 1 cursor["seed_offset"] = seed_offset cursor["seed_next_batch_ordinal"] = next_batch_ordinal cursor["seed_page_size"] = page_size cursor["seed_total_hits"] = total_hits cursor["seed_total_rows"] = total_rows cursor["seed_total_to_export"] = total_rows cursor["seed_total_batches"] = total_to_export connection.execute( """ UPDATE export_runs SET total_items = ?, cursor_json = ?, last_heartbeat_at = ? WHERE run_id = ? """, (total_to_export, compact_json_text(cursor), utc_now(), run_id), ) connection.commit() if seed_offset >= total_rows or len(page_rows) < current_page_size: completed = True break if completed: cursor["seed_offset"] = seed_offset cursor["seed_next_batch_ordinal"] = next_batch_ordinal cursor["seed_page_size"] = page_size cursor["seed_total_rows"] = total_rows cursor["seed_total_to_export"] = total_rows cursor["seed_total_batches"] = total_to_export cursor["seed_completed"] = True next_phase = "exporting" if total_to_export > 0 else "finalizing" connection.execute( """ UPDATE export_runs SET phase = ?, status = ?, total_items = ?, cursor_json = ?, last_heartbeat_at = ? WHERE run_id = ? """, (next_phase, next_phase, total_to_export, compact_json_text(cursor), utc_now(), run_id), ) connection.commit() return { "seeded": seeded_rows, "seeded_batches": seeded_batches, "offset": seed_offset, "total_rows": total_rows, "total_items": total_to_export, "completed": completed, "table": table_name, } def render_csv_fragment_for_item( connection: sqlite3.Connection, item: sqlite3.Row, field_defs: list[dict[str, str]], ) -> str: payload = decode_json_text(item["payload_json"], default={}) or {} table_name = normalize_export_table_name(str(payload.get("table") or "documents")) if table_name != "documents": rows = payload.get("rows") if isinstance(rows, list): buffer = io.StringIO() writer = csv.writer(buffer) for row in rows: if not isinstance(row, dict): continue writer.writerow( [ serialize_export_cell_value( row.get(str(field_def["field_name"])), str(field_def["field_type"]), ) for field_def in field_defs ] ) return buffer.getvalue() row = payload.get("row") if not isinstance(row, dict): raise RetrieverError(f"Missing {table_name} row payload for CSV export item {int(item['id'])}.") buffer = io.StringIO() writer = csv.writer(buffer) writer.writerow( [ serialize_export_cell_value( row.get(str(field_def["field_name"])), str(field_def["field_type"]), ) for field_def in field_defs ] ) return buffer.getvalue() document_id = int(item["document_id"]) rows = fetch_visible_document_rows_by_ids(connection, [document_id]) if not rows: raise RetrieverError(f"Document {document_id} is not visible for CSV export.") context = build_export_context(connection, rows, field_defs) row = rows[0] buffer = io.StringIO() writer = csv.writer(buffer) writer.writerow( [ serialize_export_cell_value( export_field_value(row, field_def, context), field_def["field_type"], ) for field_def in field_defs ] ) return buffer.getvalue() def export_csv_exporting_step( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, *, deadline: float, ) -> dict[str, object]: run_id = str(row["run_id"]) config = decode_json_text(row["config_json"], default={}) or {} field_defs = list(config.get("fields") or []) processed = 0 failed = 0 while ingest_v2_deadline_remaining_seconds(deadline) >= EXPORT_RUN_STEP_MIN_REMAINING_SECONDS: items = connection.execute( """ SELECT * FROM export_work_items WHERE run_id = ? AND status = 'pending' ORDER BY ordinal ASC LIMIT ? """, (run_id, EXPORT_CSV_ROW_BATCH_SIZE), ).fetchall() if not items: break for item in items: if ingest_v2_deadline_remaining_seconds(deadline) < EXPORT_RUN_STEP_MIN_REMAINING_SECONDS: break now = utc_now() try: connection.execute( "UPDATE export_work_items SET status = 'running', updated_at = ? WHERE id = ?", (now, int(item["id"])), ) fragment = render_csv_fragment_for_item(connection, item, field_defs) fragment_rel_path = export_run_csv_fragment_rel_path(run_id, int(item["ordinal"])) fragment_path = paths["state_dir"] / fragment_rel_path write_text_file_atomic(fragment_path, fragment) artifact = { "fragment_rel_path": fragment_rel_path, "bytes": len(fragment.encode("utf-8")), } now = utc_now() connection.execute( """ UPDATE export_work_items SET status = 'completed', artifact_manifest_json = ?, last_error = NULL, updated_at = ? WHERE id = ? """, (compact_json_text(artifact), now, int(item["id"])), ) processed += 1 except Exception as exc: now = utc_now() connection.execute( """ UPDATE export_work_items SET status = 'failed', last_error = ?, updated_at = ? WHERE id = ? """, (f"{type(exc).__name__}: {exc}", now, int(item["id"])), ) failed += 1 refresh_export_run_counts(connection, run_id=run_id) connection.commit() if len(items) < EXPORT_CSV_ROW_BATCH_SIZE: break pending = int(connection.execute( "SELECT COUNT(*) AS count FROM export_work_items WHERE run_id = ? AND status = 'pending'", (run_id,), ).fetchone()["count"]) running = int(connection.execute( "SELECT COUNT(*) AS count FROM export_work_items WHERE run_id = ? AND status = 'running'", (run_id,), ).fetchone()["count"]) if pending == 0 and running == 0: now = utc_now() connection.execute( """ UPDATE export_runs SET phase = 'finalizing', status = 'finalizing', last_heartbeat_at = ? WHERE run_id = ? """, (now, run_id), ) connection.commit() return {"processed": processed, "failed": failed, "pending": pending} def export_csv_finalize_step( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, *, deadline: float, ) -> dict[str, object]: run_id = str(row["run_id"]) config = decode_json_text(row["config_json"], default={}) or {} cursor = decode_json_text(row["cursor_json"], default={}) or {} field_defs = list(config.get("fields") or []) output_path = export_run_current_output_path(paths["root"], row) failed_count = int(row["failed_items"] or 0) assembled_ordinal = int(cursor.get("assembled_ordinal") or 0) partial_path = export_run_partial_csv_path(paths, run_id) if assembled_ordinal > 0 and not partial_path.exists(): assembled_ordinal = 0 completed_rows = connection.execute( """ SELECT ordinal, artifact_manifest_json FROM export_work_items WHERE run_id = ? AND status = 'completed' AND ordinal > ? ORDER BY ordinal ASC LIMIT ? """, (run_id, assembled_ordinal, EXPORT_CSV_FINALIZE_FRAGMENT_BATCH_SIZE), ).fetchall() target_ordinal = assembled_ordinal fragments: list[tuple[int, Path]] = [] for item in completed_rows: if ingest_v2_deadline_remaining_seconds(deadline) < EXPORT_RUN_STEP_MIN_REMAINING_SECONDS and fragments: break artifact = decode_json_text(item["artifact_manifest_json"], default={}) or {} fragment_rel_path = str(artifact.get("fragment_rel_path") or "") if not fragment_rel_path: continue target_ordinal = int(item["ordinal"]) fragments.append((target_ordinal, paths["state_dir"] / fragment_rel_path)) partial_path.parent.mkdir(parents=True, exist_ok=True) mode = "a" if assembled_ordinal > 0 and partial_path.exists() else "w" with partial_path.open(mode, encoding="utf-8", newline="") as handle: if mode == "w": writer = csv.writer(handle) writer.writerow([field_def["field_name"] for field_def in field_defs]) handle.flush() for _, fragment_path in fragments: if fragment_path.exists(): handle.write(fragment_path.read_text(encoding="utf-8")) cursor["assembled_ordinal"] = target_ordinal cursor["partial_rel_path"] = str(partial_path.relative_to(paths["state_dir"])) now = utc_now() connection.execute( """ UPDATE export_runs SET cursor_json = ?, last_heartbeat_at = ? WHERE run_id = ? """, (compact_json_text(cursor), now, run_id), ) connection.commit() completed_count = int(row["completed_items"] or 0) total_items = int(row["total_items"] or 0) remaining_completed = int(connection.execute( """ SELECT COUNT(*) AS count FROM export_work_items WHERE run_id = ? AND status = 'completed' AND ordinal > ? """, (run_id, target_ordinal), ).fetchone()["count"]) if completed_count + failed_count >= total_items and remaining_completed == 0: output_path.parent.mkdir(parents=True, exist_ok=True) os.replace(partial_path, output_path) status = "completed" if failed_count == 0 else "failed" error = None if failed_count == 0 else f"{failed_count} CSV export row(s) failed." connection.execute( """ UPDATE export_runs SET phase = ?, status = ?, cursor_json = ?, completed_at = ?, last_heartbeat_at = ?, error = ? WHERE run_id = ? """, (status, status, compact_json_text(cursor), now, now, error, run_id), ) connection.commit() return { "assembled_rows": target_ordinal, "assembled_this_step": len(fragments), "remaining_completed": remaining_completed, "failed": failed_count, } def archive_reset_after_corrupt_partial_zip(connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row) -> bool: run_id = str(row["run_id"]) partial_path = export_run_partial_archive_path(paths, run_id) if not partial_path.exists() or zipfile.is_zipfile(partial_path): return False partial_path.unlink() now = utc_now() cursor = decode_json_text(row["cursor_json"], default={}) or {} cursor["partial_zip_rebuilt_after_corruption_at"] = now cursor["portable_workspace_added"] = False cursor["manifest_added"] = False connection.execute( """ UPDATE export_work_items SET status = 'pending', artifact_manifest_json = '{}', last_error = NULL, updated_at = ? WHERE run_id = ? AND status = 'completed' """, (now, run_id), ) connection.execute( """ UPDATE export_runs SET phase = 'exporting', status = 'exporting', completed_items = 0, failed_items = 0, cursor_json = ?, last_heartbeat_at = ? WHERE run_id = ? """, (compact_json_text(cursor), now, run_id), ) connection.commit() return True def export_archive_exporting_step( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, *, deadline: float, ) -> dict[str, object]: run_id = str(row["run_id"]) config = decode_json_text(row["config_json"], default={}) or {} partial_path = export_run_partial_archive_path(paths, run_id) partial_path.parent.mkdir(parents=True, exist_ok=True) processed = 0 failed = 0 portable_workspace = bool(config.get("portable_workspace")) while ingest_v2_deadline_remaining_seconds(deadline) >= EXPORT_RUN_STEP_MIN_REMAINING_SECONDS: item = connection.execute( """ SELECT * FROM export_work_items WHERE run_id = ? AND status = 'pending' ORDER BY ordinal ASC LIMIT 1 """, (run_id,), ).fetchone() if item is None: break now = utc_now() connection.execute( "UPDATE export_work_items SET status = 'running', updated_at = ? WHERE id = ?", (now, int(item["id"])), ) connection.commit() try: rows = fetch_visible_document_rows_by_ids(connection, [int(item["document_id"])]) if not rows: raise RetrieverError(f"Document {int(item['document_id'])} is not visible for archive export.") payload = decode_json_text(item["payload_json"], default={}) or {} mode = "a" if partial_path.exists() else "w" with zipfile.ZipFile(partial_path, mode, compression=zipfile.ZIP_DEFLATED) as archive: written_member_paths = set(archive.namelist()) manifest_entry, document_warnings = archive_document_files( archive, written_member_paths, connection, paths, rows[0], ) manifest_entry["ordinal"] = int(item["ordinal"]) manifest_entry["inclusion_reason"] = payload.get("inclusion_reason") or {} if portable_workspace: revision_entries, revision_warnings = archive_document_text_revisions( archive, written_member_paths, connection, paths, int(item["document_id"]), ) manifest_entry["text_revision_entries"] = revision_entries document_warnings.extend(revision_warnings) artifact = { "document": manifest_entry, "warnings": document_warnings, } now = utc_now() connection.execute( """ UPDATE export_work_items SET status = 'completed', artifact_manifest_json = ?, last_error = NULL, updated_at = ? WHERE id = ? """, (compact_json_text(artifact), now, int(item["id"])), ) processed += 1 except Exception as exc: now = utc_now() connection.execute( """ UPDATE export_work_items SET status = 'failed', last_error = ?, updated_at = ? WHERE id = ? """, (f"{type(exc).__name__}: {exc}", now, int(item["id"])), ) failed += 1 refresh_export_run_counts(connection, run_id=run_id) connection.commit() pending = int(connection.execute( "SELECT COUNT(*) AS count FROM export_work_items WHERE run_id = ? AND status IN ('pending', 'running')", (run_id,), ).fetchone()["count"]) if pending == 0: now = utc_now() connection.execute( """ UPDATE export_runs SET phase = 'finalizing', status = 'finalizing', last_heartbeat_at = ? WHERE run_id = ? """, (now, run_id), ) connection.commit() return {"processed": processed, "failed": failed, "pending": pending} def export_archive_finalize_step( connection: sqlite3.Connection, paths: dict[str, Path], row: sqlite3.Row, *, deadline: float, ) -> dict[str, object]: run_id = str(row["run_id"]) if ingest_v2_deadline_remaining_seconds(deadline) < EXPORT_RUN_STEP_MIN_REMAINING_SECONDS: return {"finalized": False, "reason": "budget_exhausted"} config = decode_json_text(row["config_json"], default={}) or {} cursor = decode_json_text(row["cursor_json"], default={}) or {} selector = decode_json_text(row["selector_json"], default={}) or {} metadata_field_defs = list(config.get("metadata_fields") or []) partial_path = export_run_partial_archive_path(paths, run_id) partial_path.parent.mkdir(parents=True, exist_ok=True) output_path = export_run_current_output_path(paths["root"], row) completed_items = connection.execute( """ SELECT * FROM export_work_items WHERE run_id = ? AND status = 'completed' ORDER BY ordinal ASC """, (run_id,), ).fetchall() failed_count = int(row["failed_items"] or 0) manifest_document_entries: list[dict[str, object]] = [] warnings: list[str] = [] selected_document_ids: list[int] = [] for item in completed_items: artifact = decode_json_text(item["artifact_manifest_json"], default={}) or {} document_entry = artifact.get("document") if isinstance(artifact, dict) else None if isinstance(document_entry, dict): manifest_document_entries.append(document_entry) selected_document_ids.append(int(item["document_id"])) for warning in list(artifact.get("warnings") or []) if isinstance(artifact, dict) else []: warnings.append(str(warning)) portable_workspace_payload = None with zipfile.ZipFile(partial_path, "a" if partial_path.exists() else "w", compression=zipfile.ZIP_DEFLATED) as archive: written_member_paths = set(archive.namelist()) if bool(config.get("portable_workspace")) and not bool(cursor.get("portable_workspace_added")): rows = fetch_visible_document_rows_by_ids(connection, selected_document_ids) with tempfile.TemporaryDirectory(prefix="retriever-portable-workspace-") as tempdir: portable_root = Path(tempdir) / "workspace" portable_workspace_payload = build_portable_workspace_db(connection, portable_root, rows) add_archive_file_once( archive, written_member_paths, Path(portable_workspace_payload["db_path"]), ".retriever/retriever.db", ) cursor["portable_workspace_added"] = True cursor["portable_workspace_payload"] = portable_workspace_payload else: portable_workspace_payload = cursor.get("portable_workspace_payload") selected_rows = fetch_visible_document_rows_by_ids(connection, selected_document_ids) sidecar_payload = write_archive_sidecars( archive, written_member_paths, connection, paths, selected_rows, manifest_document_entries, metadata_payload=build_archive_metadata_csv_bytes( connection, selected_rows, manifest_document_entries, metadata_field_defs or resolve_archive_metadata_field_definitions(connection), ), created_at=str(cursor.get("created_at") or row["created_at"]), selector=selector, family_mode=str(config.get("family_mode") or ""), seed_limit=int(config["seed_limit"]) if config.get("seed_limit") is not None else None, portable_workspace=bool(config.get("portable_workspace")), portable_workspace_payload=portable_workspace_payload if isinstance(portable_workspace_payload, dict) else None, warnings=warnings, failed_count=failed_count, ) cursor["metadata_added"] = True cursor["manifest_added"] = True checksum_payload = ensure_archive_checksums_file(partial_path) output_path.parent.mkdir(parents=True, exist_ok=True) os.replace(partial_path, output_path) now = utc_now() status = "completed" if failed_count == 0 else "failed" error = None if failed_count == 0 else f"{failed_count} archive export item(s) failed." connection.execute( """ UPDATE export_runs SET phase = ?, status = ?, cursor_json = ?, completed_at = ?, last_heartbeat_at = ?, error = ? WHERE run_id = ? """, (status, status, compact_json_text(cursor), now, now, error, run_id), ) connection.commit() return { "finalized": True, "archive_member_count": int(checksum_payload["archive_member_count"]), "document_count": len(manifest_document_entries), "counts": sidecar_payload["sidecar_counts"], "warnings": len(warnings), "failed": failed_count, } def export_run_step( root: Path, *, export_kind: str, run_id: str | None = None, budget_seconds: int | None = None, ) -> dict[str, object]: budget = normalize_resumable_step_budget(budget_seconds) deadline = time.perf_counter() + max(0.1, float(budget) - 0.25) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) executed_steps: list[str] = [] step_results: list[dict[str, object]] = [] try: apply_schema(connection, root) row = latest_export_run_row(connection, export_kind=export_kind, run_id=run_id) if row is None: return { "ok": True, "executed": False, "executed_steps": [], "reason": "no_export_run", "more_work_remaining": False, "run": export_status(root, export_kind=export_kind, run_id=run_id, budget_seconds=budget), "remaining_budget_seconds": round(ingest_v2_deadline_remaining_seconds(deadline), 3), } resolved_run_id = str(row["run_id"]) connection.execute( """ UPDATE export_work_items SET status = 'pending', updated_at = ? WHERE run_id = ? AND status = 'running' """, (utc_now(), resolved_run_id), ) connection.commit() row = require_export_run_row(connection, export_kind=export_kind, run_id=resolved_run_id) if export_kind == "archive" and archive_reset_after_corrupt_partial_zip(connection, paths, row): row = require_export_run_row(connection, export_kind=export_kind, run_id=resolved_run_id) while ingest_v2_deadline_remaining_seconds(deadline) >= EXPORT_RUN_STEP_MIN_REMAINING_SECONDS: row = require_export_run_row(connection, export_kind=export_kind, run_id=resolved_run_id) phase = str(row["phase"]) status = str(row["status"]) if status in EXPORT_RUN_TERMINAL_STATUSES: break if phase == "seeding" and export_kind == "csv": step_result = export_table_csv_seed_step(root, connection, paths, row, deadline=deadline) executed_steps.append("seed") step_results.append(step_result) if int(step_result.get("seeded") or 0) == 0 and not bool(step_result.get("completed")): break continue if phase == "exporting": if export_kind == "csv": step_result = export_csv_exporting_step(connection, paths, row, deadline=deadline) else: step_result = export_archive_exporting_step(connection, paths, row, deadline=deadline) executed_steps.append("export") step_results.append(step_result) if int(step_result.get("processed") or 0) == 0 and int(step_result.get("failed") or 0) == 0: break continue if phase == "finalizing": if export_kind == "csv": step_result = export_csv_finalize_step(connection, paths, row, deadline=deadline) else: step_result = export_archive_finalize_step(connection, paths, row, deadline=deadline) executed_steps.append("finalize") step_results.append(step_result) break break updated_row = require_export_run_row(connection, export_kind=export_kind, run_id=resolved_run_id) run_payload = export_run_status_payload(connection, root, updated_row, budget_seconds=budget) reason = ( "run_terminal" if str(updated_row["status"]) in EXPORT_RUN_TERMINAL_STATUSES else "budget_exhausted" if ingest_v2_deadline_remaining_seconds(deadline) < EXPORT_RUN_STEP_MIN_REMAINING_SECONDS else "no_runnable_step" ) return { "ok": True, "executed": bool(executed_steps), "executed_steps": executed_steps, "reason": reason, "step_result": step_results[-1] if step_results else None, "step_results": step_results, "more_work_remaining": str(updated_row["status"]) not in EXPORT_RUN_TERMINAL_STATUSES, "run": run_payload, "remaining_budget_seconds": round(ingest_v2_deadline_remaining_seconds(deadline), 3), } finally: connection.close() def export_previews( root: Path, raw_output_path: str, document_ids: list[int] | None, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) output_dir = resolve_export_output_dir(paths, raw_output_path) normalized_document_ids = list(dict.fromkeys(int(document_id) for document_id in (document_ids or []))) if normalized_document_ids and (query.strip() or raw_filters or sort_field or order): raise RetrieverError("export-previews accepts either --doc-id selectors or query/filter selectors, not both.") if normalized_document_ids: rows = fetch_visible_document_rows_by_ids(connection, normalized_document_ids) selector: dict[str, object] = { "mode": "document_ids", "document_ids": normalized_document_ids, } else: selection = resolve_document_search(connection, query, raw_filters, sort_field, order) rows = [item["row"] for item in selection["results"]] selector = { "mode": "search", "query": selection["query"], "filters": selection["filters"], "sort": selection["sort"], "order": selection["order"], } output_dir.mkdir(parents=True, exist_ok=True) cleanup_previous_export_preview_outputs(output_dir) units_dir = output_dir / "units" units_dir.mkdir(parents=True, exist_ok=True) documents_dir = output_dir / "documents" documents_dir.mkdir(parents=True, exist_ok=True) units = build_export_preview_units(connection, paths, rows) export_asset_rel_paths = copy_export_preview_assets( paths, output_dir, [ dict(document) for unit in units for document in list(unit.get("documents") or []) if isinstance(document, dict) ], ) unit_payloads: list[dict[str, object]] = [] document_targets_by_id: dict[int, dict[str, object]] = {} for unit in units: file_name = export_preview_unit_file_name(unit) unit_output_path = units_dir / file_name unit_output_rel_path = Path("units") / file_name unit_output_path.write_text( build_export_preview_unit_html( unit, file_name=file_name, output_rel_path=unit_output_rel_path.as_posix(), ), encoding="utf-8", ) unit_payload = { "unit_key": str(unit["unit_key"]), "unit_kind": str(unit["unit_kind"]), "title": str(unit["title"]), "summary": str(unit["summary"]), "conversation_id": int(unit["conversation_id"]) if unit.get("conversation_id") is not None else None, "conversation_type": unit.get("conversation_type"), "document_ids": [int(document["id"]) for document in unit["documents"]], "selected_document_ids": [int(document_id) for document_id in unit["selected_document_ids"]], "output_path": str(unit_output_path), "output_rel_path": unit_output_rel_path.as_posix(), "file_size": file_size_bytes(unit_output_path), } unit_payloads.append(unit_payload) for document in unit["documents"]: document_id = int(document["id"]) if document_id not in unit_payload["selected_document_ids"]: continue target_output_path = unit_output_path target_output_rel_path = unit_output_rel_path.as_posix() target_href = unit_output_rel_path.as_posix() if str(unit["unit_kind"]) == "email_conversation" and len(unit["documents"]) > 1: document_file_name = export_preview_document_file_name(document) document_output_path = documents_dir / document_file_name document_output_rel_path = Path("documents") / document_file_name document_output_path.write_text( build_export_preview_document_html( connection, paths, unit=unit, document=document, document_output_path=document_output_path, document_output_rel_path=document_output_rel_path.as_posix(), unit_output_path=unit_output_path, ), encoding="utf-8", ) target_output_path = document_output_path target_output_rel_path = document_output_rel_path.as_posix() target_href = document_output_rel_path.as_posix() document_targets_by_id[document_id] = { "document_id": document_id, "title": conversation_preview_document_heading(document), "output_path": str(target_output_path), "output_rel_path": target_output_rel_path, "unit_output_path": str(unit_output_path), "unit_output_rel_path": unit_output_rel_path.as_posix(), "file_output_path": str(unit_output_path), "file_output_rel_path": unit_output_rel_path.as_posix(), "target_fragment": conversation_preview_anchor(document_id), "href": target_href, } index_path = output_dir / "index.html" index_path.write_text( build_export_preview_index_html( units=unit_payloads, selected_rows=rows, document_targets_by_id=document_targets_by_id, ), encoding="utf-8", ) index_rel_path = "index.html" document_targets = [ document_targets_by_id[int(row["id"])] for row in rows if int(row["id"]) in document_targets_by_id ] manifest_payload = { "status": "ok", "output_path": str(output_dir), "output_rel_path": relative_output_path_or_none(root, output_dir), "index_path": str(index_path), "index_rel_path": index_rel_path, "selected_document_count": len(rows), "unit_count": len(unit_payloads), "selector": selector, "units": unit_payloads, "document_targets": document_targets, "asset_rel_paths": export_asset_rel_paths, } manifest_path = output_dir / "manifest.json" manifest_path.write_text(json.dumps(manifest_payload, indent=2, sort_keys=True), encoding="utf-8") manifest_payload["manifest_path"] = str(manifest_path) manifest_payload["manifest_rel_path"] = "manifest.json" return manifest_payload finally: connection.close() def resolve_workspace_input_path(root: Path, raw_input_path: str) -> Path: normalized_input = raw_input_path.strip() if not normalized_input: raise RetrieverError("Input path cannot be empty.") requested_path = Path(normalized_input).expanduser() if requested_path.is_absolute(): return requested_path.resolve() return (root / requested_path).resolve() def inspect_pst_properties( root: Path, raw_pst_path: str, *, source_item_ids: list[str] | None = None, message_kind: str = "all", limit: int = 20, max_record_entries: int = 128, ) -> dict[str, object]: normalized_message_kind = normalize_whitespace(message_kind).lower() or "all" if normalized_message_kind not in {"all", "chat", "email", "calendar", "skip"}: raise RetrieverError(f"Unsupported PST message-kind filter: {message_kind!r}") if limit < 1: raise RetrieverError("PST inspection limit must be >= 1.") if max_record_entries < 1: raise RetrieverError("PST inspection max-record-entries must be >= 1.") pst_path = resolve_workspace_input_path(root, raw_pst_path) if not pst_path.exists(): raise RetrieverError(f"PST path not found: {pst_path}") if not pst_path.is_file(): raise RetrieverError(f"PST path is not a file: {pst_path}") if normalize_extension(pst_path) != PST_SOURCE_KIND: raise RetrieverError(f"Expected a .pst file, got: {pst_path.name}") normalized_source_item_ids = [ normalize_whitespace(str(source_item_id)) for source_item_id in (source_item_ids or []) if normalize_whitespace(str(source_item_id)) ] source_item_id_filter = set(normalized_source_item_ids) scanned = 0 matched = 0 messages: list[dict[str, object]] = [] for payload in iter_pst_debug_messages(pst_path, max_record_entries=max_record_entries): scanned += 1 if normalized_message_kind != "all" and str(payload["message_kind"]) != normalized_message_kind: continue if source_item_id_filter and str(payload["source_item_id"]) not in source_item_id_filter: continue matched += 1 if len(messages) >= limit: continue messages.append(payload) return { "status": "ok", "pst_path": str(pst_path), "pst_rel_path": relative_output_path_or_none(root, pst_path), "message_kind": normalized_message_kind, "source_item_ids": normalized_source_item_ids, "limit": int(limit), "max_record_entries": int(max_record_entries), "scanned": int(scanned), "matched": int(matched), "returned": len(messages), "truncated": matched > len(messages), "messages": messages, } def get_doc( root: Path, document_id: int, include_text: str, chunk_indexes: list[int] | None, ) -> dict[str, object]: paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = connection.execute( """ SELECT * FROM documents WHERE id = ? AND lifecycle_status NOT IN ('missing', 'deleted') """, (document_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown active document id: {document_id}") occurrence_row = preferred_occurrence_for_document(connection, document_id, ["o.lifecycle_status = 'active'"], []) chunk_rows = document_chunk_rows(connection, document_id) requested_chunk_indexes = sorted(dict.fromkeys(int(chunk_index) for chunk_index in (chunk_indexes or []))) if len(requested_chunk_indexes) > MAX_GET_DOC_CHUNKS: raise RetrieverError(f"Requested too many chunks; max is {MAX_GET_DOC_CHUNKS}.") chunk_rows_by_index = {int(chunk_row["chunk_index"]): chunk_row for chunk_row in chunk_rows} missing_chunk_indexes = [chunk_index for chunk_index in requested_chunk_indexes if chunk_index not in chunk_rows_by_index] if missing_chunk_indexes: raise RetrieverError( f"Unknown chunk indexes for document {document_id}: {', '.join(str(chunk_index) for chunk_index in missing_chunk_indexes)}" ) exact_chunks: list[dict[str, object]] = [] total_text_chars = 0 document_payload = document_overview_payload( paths, connection, row, occurrence_row=occurrence_row, ) preview_rel_path = str(document_payload["preview_rel_path"]) preview_abs_path = str(document_payload["preview_abs_path"]) for chunk_index in requested_chunk_indexes: chunk_row = chunk_rows_by_index[chunk_index] text_content = str(chunk_row["text_content"] or "") total_text_chars += len(text_content) if total_text_chars > MAX_GET_DOC_TEXT_CHARS: raise RetrieverError(f"Requested chunk text exceeds the {MAX_GET_DOC_TEXT_CHARS}-character limit.") snippet = chunk_preview_text(text_content) exact_chunks.append( { "chunk_index": int(chunk_row["chunk_index"]), "char_start": int(chunk_row["char_start"]), "char_end": int(chunk_row["char_end"]), "token_estimate": chunk_row["token_estimate"], "text": text_content, "snippet": snippet, "citation": build_chunk_citation_payload( row, occurrence_row, preview_rel_path=preview_rel_path, preview_abs_path=preview_abs_path, chunk_index=int(chunk_row["chunk_index"]), char_start=int(chunk_row["char_start"]), char_end=int(chunk_row["char_end"]), snippet=snippet, ), } ) text_summary = None normalized_include_text = include_text.strip().lower() if normalized_include_text == "summary": text_summary = reconstruct_document_text_prefix(chunk_rows, GET_DOC_SUMMARY_CHARS) elif normalized_include_text != "none": raise RetrieverError(f"Unsupported include-text mode: {include_text}") return { "status": "ok", "document": document_payload, "chunk_count": len(chunk_rows), "include_text": normalized_include_text, "text_summary": text_summary, "chunks": exact_chunks, } finally: connection.close() def list_chunks(root: Path, document_id: int, page: int, per_page: int) -> dict[str, object]: if page < 1: raise RetrieverError("Page must be >= 1.") if per_page < 1: raise RetrieverError("per-page must be >= 1.") per_page = min(per_page, MAX_CHUNK_PAGE_SIZE) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) row = connection.execute( """ SELECT * FROM documents WHERE id = ? AND lifecycle_status NOT IN ('missing', 'deleted') """, (document_id,), ).fetchone() if row is None: raise RetrieverError(f"Unknown active document id: {document_id}") occurrence_row = preferred_occurrence_for_document(connection, document_id, ["o.lifecycle_status = 'active'"], []) chunk_rows = document_chunk_rows(connection, document_id) total_chunks = len(chunk_rows) total_pages = max(1, (total_chunks + per_page - 1) // per_page) start = (page - 1) * per_page end = start + per_page paged_rows = chunk_rows[start:end] return { "status": "ok", "document": { "document_id": int(row["id"]), "control_number": row["control_number"], "file_name": occurrence_row["file_name"] if occurrence_row is not None else row["file_name"], }, "page": page, "per_page": per_page, "total_chunks": total_chunks, "total_pages": total_pages, "chunks": [ { "chunk_index": int(chunk_row["chunk_index"]), "char_start": int(chunk_row["char_start"]), "char_end": int(chunk_row["char_end"]), "token_estimate": chunk_row["token_estimate"], "snippet": chunk_preview_text(chunk_row["text_content"]), } for chunk_row in paged_rows ], } finally: connection.close() def search_chunk_rows( connection: sqlite3.Connection, query: str, clauses: list[str], params: list[object], ) -> list[sqlite3.Row]: query_value = query.strip() if not query_value: raise RetrieverError("search-chunks requires a non-empty query.") where_clause = " AND ".join(clauses) sql = f""" SELECT d.*, dc.id AS chunk_row_id, dc.chunk_index, dc.char_start, dc.char_end, dc.token_estimate, dc.text_content, bm25(chunks_fts) AS rank FROM chunks_fts JOIN document_chunks dc ON dc.id = CAST(chunks_fts.chunk_id AS INTEGER) JOIN documents d ON d.id = dc.document_id WHERE chunks_fts MATCH ? AND {where_clause} """ try: return connection.execute(sql, [query_value, *params]).fetchall() except sqlite3.OperationalError: return connection.execute(sql, [f'"{query_value}"', *params]).fetchall() def sort_chunk_match_rows(rows: list[sqlite3.Row], sort_field: str | None, order: str | None) -> list[sqlite3.Row]: normalized_sort_field = (sort_field or "relevance").lower() normalized_order = (order or ("asc" if normalized_sort_field == "relevance" else "desc")).lower() if normalized_sort_field not in {"relevance", "date_created", "date_modified"}: raise RetrieverError(f"Unsupported chunk sort field: {sort_field}") stable_rows = sorted(rows, key=lambda row: (int(row["id"]), int(row["chunk_index"]))) if normalized_sort_field == "relevance": return sorted(stable_rows, key=lambda row: float(row["rank"]), reverse=normalized_order == "desc") ranked_rows = sorted(stable_rows, key=lambda row: float(row["rank"])) return sorted(ranked_rows, key=lambda row: coerce_sort_value(row[normalized_sort_field]), reverse=normalized_order == "desc") def count_distinct_chunk_documents( connection: sqlite3.Connection, query: str, clauses: list[str], params: list[object], ) -> int: query_value = query.strip() where_clause = " AND ".join(clauses) sql = f""" SELECT COUNT(DISTINCT d.id) AS count FROM chunks_fts JOIN document_chunks dc ON dc.id = CAST(chunks_fts.chunk_id AS INTEGER) JOIN documents d ON d.id = dc.document_id WHERE chunks_fts MATCH ? AND {where_clause} """ try: row = connection.execute(sql, [query_value, *params]).fetchone() except sqlite3.OperationalError: row = connection.execute(sql, [f'"{query_value}"', *params]).fetchone() return int(row["count"] or 0) def count_filtered_documents( connection: sqlite3.Connection, clauses: list[str], params: list[object], ) -> int: row = connection.execute( f""" SELECT COUNT(*) AS count FROM documents d WHERE {' AND '.join(clauses)} """, params, ).fetchone() return int(row["count"] or 0) def build_analysis_scope_filters( connection: sqlite3.Connection, paths: dict[str, Path], *, raw_filters: list[list[str]] | None, select_from_scope: bool, ) -> tuple[dict[str, object] | None, list[str], list[object], list[object]]: if not select_from_scope: filter_summary, clauses, params = build_search_filters(connection, raw_filters) return None, clauses, params, filter_summary merged_scope = merge_scope_with_search_inputs(read_session_state(paths).get("scope"), "", raw_filters) scope, clauses, params, filter_summary = build_scope_search_filters(connection, merged_scope) keyword_query = normalize_inline_whitespace(str(scope.get("keyword") or "")) bates_query = format_scope_bates_value(scope.get("bates")) if not keyword_query and not bates_query: return scope, clauses, params, filter_summary selection = resolve_scope_document_search(connection, scope) selected_document_ids = [int(item["id"]) for item in selection["results"]] if not selected_document_ids: return selection["scope"], ["0"], [], selection["filters"] connection.execute("DROP TABLE IF EXISTS temp_scope_selected_documents") connection.execute( """ CREATE TEMP TABLE temp_scope_selected_documents ( document_id INTEGER PRIMARY KEY ) """ ) connection.executemany( "INSERT INTO temp_scope_selected_documents (document_id) VALUES (?)", [(document_id,) for document_id in selected_document_ids], ) return ( selection["scope"], ["EXISTS (SELECT 1 FROM temp_scope_selected_documents tsd WHERE tsd.document_id = d.id)"], [], selection["filters"], ) def search_chunks( root: Path, query: str, raw_filters: list[list[str]] | None, sort_field: str | None, order: str | None, top_k: int, per_doc_cap: int, *, select_from_scope: bool = False, count_only: bool = False, distinct_docs: bool = False, ) -> dict[str, object]: if top_k < 1: raise RetrieverError("top-k must be >= 1.") if per_doc_cap < 1: raise RetrieverError("per-doc-cap must be >= 1.") top_k = min(top_k, MAX_CHUNK_SEARCH_TOP_K) per_doc_cap = min(per_doc_cap, MAX_CHUNK_SEARCH_PER_DOC_CAP) if distinct_docs and not count_only: raise RetrieverError("--distinct-docs requires --count-only.") if count_only and not distinct_docs: raise RetrieverError("--count-only currently requires --distinct-docs.") paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) effective_scope, clauses, params, filter_summary = build_analysis_scope_filters( connection, paths, raw_filters=raw_filters, select_from_scope=select_from_scope, ) occurrence_scope_clauses, occurrence_scope_params = build_occurrence_scope_filters(connection, raw_filters) if count_only: payload = { "query": query, "filters": filter_summary, "documents_with_hits": count_distinct_chunk_documents(connection, query, clauses, params), "total_docs_filtered": count_filtered_documents(connection, clauses, params), "count_mode": "distinct-documents", } if effective_scope is not None: payload["selected_from_scope"] = True payload["scope"] = effective_scope return payload raw_rows = search_chunk_rows(connection, query, clauses, params) grouped_rows: dict[int, list[sqlite3.Row]] = defaultdict(list) for row in sorted(raw_rows, key=lambda item: (float(item["rank"]), int(item["chunk_index"]))): grouped_rows[int(row["id"])].append(row) selected_rows: list[sqlite3.Row] = [] for rows in grouped_rows.values(): selected_rows.extend(rows[:per_doc_cap]) sorted_rows = sort_chunk_match_rows(selected_rows, sort_field, order) returned_rows: list[sqlite3.Row] = [] total_text_chars = 0 for row in sorted_rows: text_content = str(row["text_content"] or "") if len(returned_rows) >= top_k: break if returned_rows and total_text_chars + len(text_content) > MAX_CHUNK_SEARCH_TEXT_CHARS: break total_text_chars += len(text_content) returned_rows.append(row) results: list[dict[str, object]] = [] dataset_memberships = fetch_document_dataset_memberships(connection, returned_rows) production_names = fetch_production_names(connection, returned_rows) preferred_occurrences = preferred_occurrences_by_document( connection, [int(row["id"]) for row in returned_rows], occurrence_scope_clauses, occurrence_scope_params, ) parent_summaries = fetch_parent_summaries( connection, [row for row in returned_rows if row["parent_document_id"] is not None], ) for row in returned_rows: occurrence_row = preferred_occurrences.get(int(row["id"])) path_payload = document_path_payload( paths, connection, row, occurrence_row=occurrence_row, ) preview_rel_path = str(path_payload["preview_rel_path"]) preview_abs_path = str(path_payload["preview_abs_path"]) snippet = make_snippet(str(row["text_content"] or ""), query) source_row = occurrence_row or row custodian_values = document_custodian_values_from_row(row) custodian_text = ", ".join(custodian_values) if custodian_values else None result = { **path_payload, "document_id": int(row["id"]), "control_number": row["control_number"], "file_name": source_row["file_name"], "file_type": source_row["file_type"], "custodian": custodian_text, "custodians": custodian_values, "chunk_index": int(row["chunk_index"]), "char_start": int(row["char_start"]), "char_end": int(row["char_end"]), "token_estimate": row["token_estimate"], "text": str(row["text_content"] or ""), "snippet": snippet, "rank": float(row["rank"]), "metadata": { "author": row["author"], "content_type": row["content_type"], "custodian": custodian_text, "custodians": custodian_values, "date_created": row["date_created"], "date_modified": row["date_modified"], "participants": row["participants"], "recipients": row["recipients"], "subject": row["subject"], "title": row["title"], "updated_at": row["updated_at"], }, "citation": build_chunk_citation_payload( row, occurrence_row, preview_rel_path=preview_rel_path, preview_abs_path=preview_abs_path, chunk_index=int(row["chunk_index"]), char_start=int(row["char_start"]), char_end=int(row["char_end"]), snippet=snippet, ), } membership = dataset_memberships.get(int(row["id"]), {"ids": [], "names": []}) dataset_ids = [int(dataset_id) for dataset_id in membership["ids"]] dataset_names = [str(dataset_name) for dataset_name in membership["names"]] result["dataset_ids"] = dataset_ids result["dataset_names"] = dataset_names if len(dataset_ids) == 1: result["dataset_id"] = dataset_ids[0] if len(dataset_names) == 1: result["dataset_name"] = dataset_names[0] if source_row["production_id"] is not None: result["production_name"] = production_names.get(int(source_row["production_id"])) if row["parent_document_id"] is not None: result["parent"] = parent_summaries.get(int(row["parent_document_id"])) results.append(result) normalized_sort = (sort_field or "relevance").lower() normalized_order = (order or ("asc" if normalized_sort == "relevance" else "desc")).lower() payload = { "query": query, "filters": filter_summary, "sort": normalized_sort, "order": normalized_order, "top_k": top_k, "per_doc_cap": per_doc_cap, "total_matches": len(raw_rows), "results": results, } if effective_scope is not None: payload["selected_from_scope"] = True payload["scope"] = effective_scope return payload finally: connection.close() ENTITY_AGGREGATE_GROUPS: dict[str, dict[str, object]] = { "entity_type": { "output_name": "entity_type", "select_sql": "e.entity_type", "join_documents": False, }, "entity_origin": { "output_name": "entity_origin", "select_sql": "e.entity_origin", "join_documents": False, }, "canonical_status": { "output_name": "entity_status", "select_sql": "e.canonical_status", "join_documents": False, }, "role": { "output_name": "entity_role", "select_sql": "de.role", "join_documents": True, }, } ENTITY_AGGREGATE_GROUP_ALIASES = { "entity_status": "canonical_status", "entity_role": "role", } def aggregate_output_name(base_name: str, used_names: set[str]) -> str: candidate = base_name if candidate not in used_names: used_names.add(candidate) return candidate suffix = 2 while f"{base_name}_{suffix}" in used_names: suffix += 1 candidate = f"{base_name}_{suffix}" used_names.add(candidate) return candidate def aggregate_date_base_expr(field_name: str) -> str: return f"substr(COALESCE(d.{quote_identifier(field_name)}, ''), 1, 10)" def normalize_entity_aggregate_group(raw_group: str) -> str | None: token = normalize_whitespace(raw_group) canonical = ENTITY_AGGREGATE_GROUP_ALIASES.get(token) if canonical is None and token in {"entity_type", "entity_origin"}: canonical = token if canonical in ENTITY_AGGREGATE_GROUPS: return canonical return None def resolve_entity_aggregate_group(raw_group: str, used_names: set[str]) -> dict[str, object]: token = normalize_whitespace(raw_group) canonical = normalize_entity_aggregate_group(token) if canonical is None: raise RetrieverError(f"Unsupported entity aggregate group: {raw_group}") spec = ENTITY_AGGREGATE_GROUPS[canonical] output_name = aggregate_output_name(str(spec["output_name"]), used_names) return { "token": token, "normalized_token": canonical, "output_name": output_name, "select_sql": str(spec["select_sql"]), "group_sql": str(spec["select_sql"]), "join_documents": bool(spec["join_documents"]), "is_temporal": False, } def raw_group_bys_are_entity_aggregate_groups(raw_group_bys: list[str]) -> bool: return any(normalize_entity_aggregate_group(raw_group) is not None for raw_group in raw_group_bys) def aggregate_entities( root: Path, raw_filters: list[list[str]] | None, raw_group_bys: list[str], order_by: str | None, order: str | None, limit: int, explain: bool, *, select_from_scope: bool = False, ) -> dict[str, object]: entity_group_flags = [normalize_entity_aggregate_group(raw_group) is not None for raw_group in raw_group_bys] if not all(entity_group_flags): raise RetrieverError( "Entity aggregate groups cannot be mixed with document aggregate groups. " "Use entity_type, entity_origin, entity_status, or entity_role." ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) used_names: set[str] = set() group_defs = [resolve_entity_aggregate_group(raw_group, used_names) for raw_group in raw_group_bys] join_documents = bool(raw_filters) or bool(select_from_scope) or any( bool(group_def["join_documents"]) for group_def in group_defs ) effective_scope: dict[str, object] | None = None filter_summary: list[object] = [] clauses: list[str] = [] params: list[object] = [] joins: list[str] = [] if join_documents: effective_scope, document_clauses, document_params, filter_summary = build_analysis_scope_filters( connection, paths, raw_filters=raw_filters, select_from_scope=select_from_scope, ) joins = [ "JOIN document_entities de ON de.entity_id = e.id", "JOIN documents d ON d.id = de.document_id", ] clauses.extend(document_clauses) params.extend(document_params) elif raw_filters: raise RetrieverError("Entity aggregate filters require document-scoped entity links.") group_tokens = {str(group_def["normalized_token"]) for group_def in group_defs} if "canonical_status" not in group_tokens: clauses.append("e.canonical_status = ?") params.append(ENTITY_STATUS_ACTIVE) where_sql = " AND ".join(clauses) if clauses else "1" select_parts = [ f"{group_def['select_sql']} AS {quote_identifier(str(group_def['output_name']))}" for group_def in group_defs ] select_parts.append("COUNT(DISTINCT e.id) AS count") group_parts = [str(group_def["group_sql"]) for group_def in group_defs] sql = ( f"SELECT {', '.join(select_parts)} " f"FROM entities e {' '.join(joins)} " f"WHERE {where_sql} " f"GROUP BY {', '.join(group_parts)}" ) order_name_map = { "metric": "count", **{str(group_def["token"]): str(group_def["output_name"]) for group_def in group_defs}, **{str(group_def["normalized_token"]): str(group_def["output_name"]) for group_def in group_defs}, **{str(group_def["output_name"]): str(group_def["output_name"]) for group_def in group_defs}, } resolved_order_by = order_name_map.get(order_by or "", None) if order_by else "count" if resolved_order_by is None: raise RetrieverError(f"Unsupported aggregate order-by: {order_by}") normalized_order = (order or ("desc" if resolved_order_by == "count" else "asc")).lower() sql += f" ORDER BY {quote_identifier(resolved_order_by)} {normalized_order.upper()}" for group_def in group_defs: if str(group_def["output_name"]) == resolved_order_by: continue sql += f", {quote_identifier(str(group_def['output_name']))} ASC" sql += " LIMIT ?" rows = connection.execute(sql, [*params, limit]).fetchall() buckets = [] for row in rows: bucket = {str(group_def["output_name"]): row[str(group_def["output_name"])] for group_def in group_defs} bucket["count"] = int(row["count"] or 0) buckets.append(bucket) payload = { "filters": filter_summary, "metric": "count", "aggregate_scope": "entities", "group_by": [str(group_def["token"]) for group_def in group_defs], "buckets": buckets, "graph": graph_metadata_for_buckets(group_defs, buckets), } if effective_scope is not None: payload["selected_from_scope"] = True payload["scope"] = effective_scope if explain: payload["sql"] = sql.replace(" ?", f" {limit}") return payload finally: connection.close() def resolve_aggregate_group( connection: sqlite3.Connection, raw_group: str, used_names: set[str], ) -> dict[str, object]: token = normalize_whitespace(raw_group) if not token: raise RetrieverError("Empty group-by expression.") temporal_match = re.fullmatch(r"(year|quarter|month|week):([A-Za-z0-9_]+)", token) if temporal_match: granularity = temporal_match.group(1) field_name = temporal_match.group(2) field_def = resolve_field_definition(connection, field_name) if field_def["field_type"] != "date": raise RetrieverError(f"Date bucket '{token}' requires a date-typed field.") if field_def.get("source") == "virtual": raise RetrieverError(f"Date bucket '{token}' cannot target a virtual field.") date_expr = aggregate_date_base_expr(field_def["field_name"]) if granularity == "year": select_sql = f"NULLIF(substr({date_expr}, 1, 4), '')" elif granularity == "month": select_sql = f"NULLIF(substr({date_expr}, 1, 7), '')" elif granularity == "quarter": select_sql = ( f"CASE WHEN length({date_expr}) >= 7 THEN " f"substr({date_expr}, 1, 4) || '-Q' || " f"CASE " f"WHEN CAST(substr({date_expr}, 6, 2) AS INTEGER) BETWEEN 1 AND 3 THEN '1' " f"WHEN CAST(substr({date_expr}, 6, 2) AS INTEGER) BETWEEN 4 AND 6 THEN '2' " f"WHEN CAST(substr({date_expr}, 6, 2) AS INTEGER) BETWEEN 7 AND 9 THEN '3' " f"WHEN CAST(substr({date_expr}, 6, 2) AS INTEGER) BETWEEN 10 AND 12 THEN '4' " f"ELSE NULL END " f"ELSE NULL END" ) else: select_sql = f"CASE WHEN length({date_expr}) = 10 THEN strftime('%Y-W%W', {date_expr}) ELSE NULL END" output_name = aggregate_output_name(granularity, used_names) return { "token": token, "output_name": output_name, "select_sql": select_sql, "group_sql": select_sql, "join_dataset": False, "is_temporal": True, } field_def = resolve_field_definition(connection, token) if field_def.get("source") == "virtual": if field_def["field_name"] not in AGGREGATABLE_VIRTUAL_FIELDS: raise RetrieverError(f"Field '{token}' is not aggregatable.") if field_def["field_name"] == "dataset_name": output_name = aggregate_output_name("dataset_name", used_names) return { "token": token, "output_name": output_name, "select_sql": "ds.dataset_name", "group_sql": "ds.dataset_name", "join_dataset": True, "is_temporal": False, } raise RetrieverError(f"Unsupported aggregatable virtual field: {token}") column_expr = f"d.{quote_identifier(field_def['field_name'])}" output_name = aggregate_output_name(field_def["field_name"], used_names) return { "token": token, "output_name": output_name, "select_sql": column_expr, "group_sql": column_expr, "join_dataset": False, "is_temporal": False, } def graph_metadata_for_buckets(group_defs: list[dict[str, object]], buckets: list[dict[str, object]]) -> dict[str, object]: graph_type = "bar" if any(bool(group_def["is_temporal"]) for group_def in group_defs): graph_type = "line" elif len(group_defs) == 1 and len(buckets) <= 6: graph_type = "pie" description = "Count by " + ", ".join(passive_field_label(group_def["output_name"]) for group_def in group_defs) graph: dict[str, object] = { "type": graph_type, "x_axis": group_defs[0]["output_name"] if group_defs else None, "y_axis": "count", "description": description, } if len(group_defs) > 1: graph["series"] = group_defs[1]["output_name"] return graph def aggregate( root: Path, raw_filters: list[list[str]] | None, raw_group_bys: list[str], metric: str, order_by: str | None, order: str | None, limit: int, explain: bool, *, select_from_scope: bool = False, ) -> dict[str, object]: if not raw_group_bys: raise RetrieverError("aggregate requires at least one --group-by.") if metric.strip().lower() != "count": raise RetrieverError("aggregate currently supports only metric=count.") if limit < 1: raise RetrieverError("limit must be >= 1.") limit = min(limit, MAX_AGGREGATE_LIMIT) if raw_group_bys_are_entity_aggregate_groups(raw_group_bys): return aggregate_entities( root, raw_filters, raw_group_bys, order_by, order, limit, explain, select_from_scope=select_from_scope, ) paths = workspace_paths(root) ensure_layout(paths) connection = connect_db(paths["db_path"]) try: apply_schema(connection, root) effective_scope, clauses, params, filter_summary = build_analysis_scope_filters( connection, paths, raw_filters=raw_filters, select_from_scope=select_from_scope, ) used_names: set[str] = set() group_defs = [resolve_aggregate_group(connection, raw_group, used_names) for raw_group in raw_group_bys] joins: list[str] = [] if any(bool(group_def["join_dataset"]) for group_def in group_defs): joins.extend( [ "JOIN dataset_documents dd ON dd.document_id = d.id", "JOIN datasets ds ON ds.id = dd.dataset_id", ] ) select_parts = [f"{group_def['select_sql']} AS {quote_identifier(str(group_def['output_name']))}" for group_def in group_defs] select_parts.append(("COUNT(DISTINCT d.id)" if joins else "COUNT(*)") + " AS count") group_parts = [str(group_def["group_sql"]) for group_def in group_defs] sql = ( f"SELECT {', '.join(select_parts)} " f"FROM documents d {' '.join(joins)} " f"WHERE {' AND '.join(clauses)} " f"GROUP BY {', '.join(group_parts)}" ) order_name_map = { "metric": "count", **{str(group_def["token"]): str(group_def["output_name"]) for group_def in group_defs}, **{str(group_def["output_name"]): str(group_def["output_name"]) for group_def in group_defs}, } resolved_order_by = order_name_map.get(order_by or "", None) if order_by else "count" if resolved_order_by is None: raise RetrieverError(f"Unsupported aggregate order-by: {order_by}") normalized_order = (order or ("desc" if resolved_order_by == "count" else "asc")).lower() sql += f" ORDER BY {quote_identifier(resolved_order_by)} {normalized_order.upper()}" if group_defs: for group_def in group_defs: if str(group_def["output_name"]) == resolved_order_by: continue sql += f", {quote_identifier(str(group_def['output_name']))} ASC" sql += " LIMIT ?" rows = connection.execute(sql, [*params, limit]).fetchall() buckets = [] for row in rows: bucket = {str(group_def["output_name"]): row[str(group_def["output_name"])] for group_def in group_defs} bucket["count"] = int(row["count"] or 0) buckets.append(bucket) payload = { "filters": filter_summary, "metric": "count", "group_by": [str(group_def["token"]) for group_def in group_defs], "buckets": buckets, "graph": graph_metadata_for_buckets(group_defs, buckets), } if effective_scope is not None: payload["selected_from_scope"] = True payload["scope"] = effective_scope if explain: payload["sql"] = sql.replace(" ?", f" {limit}") return payload finally: connection.close() def add_dataset_selector_arguments(parser: argparse.ArgumentParser) -> None: selector_group = parser.add_mutually_exclusive_group(required=True) selector_group.add_argument("--dataset-id", type=int, help="Dataset id") selector_group.add_argument("--dataset-name", help="Exact dataset name") DATASET_POLICY_BOOLEAN_CHOICES = ("true", "false", "yes", "no", "1", "0", "on", "off", "enabled", "disabled") def add_dataset_policy_boolean_argument( parser: argparse.ArgumentParser, flag: str, *, dest: str, help_text: str, ) -> None: parser.add_argument(flag, dest=dest, choices=DATASET_POLICY_BOOLEAN_CHOICES, help=help_text) def add_search_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("workspace", help="Workspace root path") parser.add_argument("query", nargs="?", default="", help="Keyword query text") parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) parser.add_argument("--sort", "--sort-by", dest="sort", help="Sort field or 'relevance'") parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") parser.add_argument("--page", type=int, default=1, help="1-based result page") parser.add_argument( "--per-page", "--limit", dest="per_page", type=int, default=None, help="Results per page (defaults to saved /page-size or 10)", ) parser.add_argument("--columns", help="Comma-separated result columns") parser.add_argument("--mode", choices=("compose", "view"), default="compose", help="Search response mode") parser.add_argument( "--verbose", action="store_true", help="Return the full payload instead of the default compact JSON", ) def add_scope_run_selector_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("--keyword", dest="query", help="Keyword query text") parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) parser.add_argument("--bates", help="Bates token or Bates range") parser.add_argument( "--dataset", dest="dataset_names", action="append", help="Exact dataset name (repeatable)", ) parser.add_argument("--from-run-id", type=int, help="Restrict to documents already present in a prior run") parser.add_argument( "--select-from-scope", action="store_true", help="AND-narrow the selector with the persisted workspace scope", ) parser.add_argument("query", nargs="?", default="", help="Optional keyword query text") def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Retriever workspace tool") parser.add_argument( "--output", choices=("json", "human"), default="json", help="Output mode. Use `human` for concise terminal-friendly summaries when supported.", ) parser.add_argument( "--human", action="store_true", help="Shortcut for `--output human`.", ) subparsers = parser.add_subparsers(dest="command", required=True) workspace_parser = subparsers.add_parser( "workspace", help="Initialize, inspect, or update workspace installation and schema", ) workspace_subparsers = workspace_parser.add_subparsers(dest="workspace_action", required=True) workspace_init_parser = workspace_subparsers.add_parser( "init", help="Initialize or repair workspace schema state and runtime metadata", ) workspace_init_parser.add_argument("workspace", help="Workspace root path") workspace_init_parser.add_argument( "--quick", action="store_true", help="Return the compact status report after initialization", ) workspace_status_parser = workspace_subparsers.add_parser( "status", help="Check runtime and workspace readiness without refreshing runtime metadata", ) workspace_status_parser.add_argument("workspace", help="Workspace root path") workspace_status_parser.add_argument( "--quick", action="store_true", help="Return the compact runtime payload", ) workspace_update_parser = workspace_subparsers.add_parser( "update", help="Refresh workspace runtime metadata from the canonical tools.py bundle", ) workspace_update_parser.add_argument("workspace", help="Workspace root path") workspace_update_parser.add_argument( "--from", dest="canonical_source", default=None, help="Path to the canonical tools.py bundle (defaults to auto-discovery)", ) workspace_update_parser.add_argument( "--force", action="store_true", help="Ignored compatibility flag retained for older callers", ) workspace_doctor_parser = workspace_subparsers.add_parser( "doctor", help="Run deeper workspace and database diagnostics", ) workspace_doctor_parser.add_argument("workspace", help="Workspace root path") workspace_doctor_parser.add_argument( "--quick", action="store_true", help="Return the compact diagnostic payload", ) workspace_doctor_parser.add_argument( "--repair-stale-sidecars", action="store_true", help="Remove stale SQLite sidecars before opening the database", ) ingest_parser = subparsers.add_parser("ingest", help="Index documents in the workspace") ingest_parser.add_argument("workspace", help="Workspace root path") ingest_parser.add_argument("--recursive", action="store_true", help="Scan directories recursively") ingest_parser.add_argument( "--path", dest="paths", action="append", help="Limit ingest to a file or directory inside the workspace; repeat to scan multiple paths", ) ingest_parser.add_argument( "--file-types", help="Comma-separated file types to include, e.g. pdf,docx,eml", ) ingest_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Per-call budget for the bounded resumable ingest facade", ) ingest_parser.add_argument( "--run-to-completion", action="store_true", help="Keep advancing resumable ingest steps until the run reaches a terminal state", ) ingest_start_parser = subparsers.add_parser("ingest-start", help="Start a resumable V2 ingest run") ingest_start_parser.add_argument("workspace", help="Workspace root path") ingest_start_parser.add_argument("--recursive", action="store_true", help="Scan directories recursively") ingest_start_parser.add_argument( "--path", dest="paths", action="append", help="Limit ingest to a file or directory inside the workspace; repeat to scan multiple paths", ) ingest_start_parser.add_argument( "--file-types", help="Comma-separated file types to include, e.g. pdf,docx,eml", ) ingest_start_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Recommended per-step budget; values above the bounded-worker cap are rejected", ) ingest_status_parser = subparsers.add_parser("ingest-status", help="Show resumable V2 ingest status") ingest_status_parser.add_argument("workspace", help="Workspace root path") ingest_status_parser.add_argument("--run-id", help="Run id; defaults to the latest resumable ingest run") ingest_status_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Budget used when computing next recommended commands", ) ingest_cancel_parser = subparsers.add_parser("ingest-cancel", help="Cancel a resumable V2 ingest run") ingest_cancel_parser.add_argument("workspace", help="Workspace root path") ingest_cancel_parser.add_argument("--run-id", required=True, help="Run id to cancel") ingest_cancel_parser.add_argument("--force", action="store_true", help="Compatibility flag reserved for future workers") ingest_run_step_parser = subparsers.add_parser( "ingest-run-step", help="Run recommended resumable V2 ingest steps until the bounded call budget is nearly exhausted", ) ingest_run_step_parser.add_argument("workspace", help="Workspace root path") ingest_run_step_parser.add_argument("--run-id", help="Run id to advance; defaults to the latest run") ingest_run_step_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Per-call budget; values above the bounded-worker cap are rejected", ) for step_command, help_text in ( ("ingest-plan-step", "Advance resumable V2 ingest planning"), ("ingest-prepare-step", "Prepare resumable V2 ingest work items"), ("ingest-commit-step", "Commit prepared resumable V2 ingest work items"), ("ingest-finalize-step", "Advance resumable V2 ingest finalization"), ): step_parser = subparsers.add_parser(step_command, help=help_text) step_parser.add_argument("workspace", help="Workspace root path") step_parser.add_argument("--run-id", required=True, help="Run id to advance") step_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Per-call budget; values above the bounded-worker cap are rejected", ) if step_command == "ingest-commit-step": step_parser.add_argument( "--max-items", type=int, help="Maximum prepared items to commit in this call", ) ingest_production_parser = subparsers.add_parser("ingest-production", help="Ingest a processed production volume") ingest_production_parser.add_argument("workspace", help="Workspace root path") ingest_production_parser.add_argument("production_root", help="Production root directory inside the workspace") inspect_pst_parser = subparsers.add_parser( "inspect-pst-properties", help="Inspect raw PST message fields and named-property candidates for debugging conversation scope ids", ) inspect_pst_parser.add_argument("workspace", help="Workspace root path") inspect_pst_parser.add_argument("pst_path", help="PST file path; relative paths resolve from the workspace root") inspect_pst_parser.add_argument( "--source-item-id", dest="source_item_ids", action="append", help="Restrict to one PST message source_item_id (repeatable)", ) inspect_pst_parser.add_argument( "--message-kind", default="all", choices=("all", "chat", "email", "calendar", "skip"), help="Filter to one normalized PST message kind", ) inspect_pst_parser.add_argument( "--limit", type=int, default=20, help="Maximum number of matching messages to return", ) inspect_pst_parser.add_argument( "--max-record-entries", type=int, default=128, help="Maximum number of record entries to serialize per record set", ) search_parser = subparsers.add_parser("search", help="Search indexed documents") add_search_arguments(search_parser) slash_parser = subparsers.add_parser("slash", help="Execute a scope-aware slash command") slash_parser.add_argument("workspace", help="Workspace root path") slash_parser.add_argument("command_text", nargs=argparse.REMAINDER, help="Slash command text") catalog_parser = subparsers.add_parser("catalog", help="Describe searchable, filterable, and aggregatable fields") catalog_parser.add_argument("workspace", help="Workspace root path") export_parser = subparsers.add_parser("export-csv", help="Write selected documents and fields to a CSV on disk") export_parser.add_argument("workspace", help="Workspace root path") export_parser.add_argument("output_path", help="CSV file path; relative paths resolve under .retriever/exports") export_parser.add_argument("query", nargs="?", default="", help="Optional keyword query text for search-based export") export_parser.add_argument( "--field", dest="fields", action="append", required=True, help="Field to export (repeatable, preserves order)", ) export_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to export (repeatable, preserves input order)", ) export_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) export_parser.add_argument( "--select-from-scope", action="store_true", help="AND-narrow the export selector with the persisted workspace scope", ) export_parser.add_argument("--sort", "--sort-by", dest="sort", help="Sort field for search-based export or 'relevance'") export_parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") export_csv_start_parser = subparsers.add_parser( "export-csv-start", help="Start a resumable CSV export", ) export_csv_start_parser.add_argument("workspace", help="Workspace root path") export_csv_start_parser.add_argument("output_path", help="CSV file path; relative paths resolve under .retriever/exports") export_csv_start_parser.add_argument("query", nargs="?", default="", help="Optional keyword query text for search-based export") export_csv_start_parser.add_argument( "--field", dest="fields", action="append", required=True, help="Field to export (repeatable, preserves order)", ) export_csv_start_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to export (repeatable, preserves input order)", ) export_csv_start_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) export_csv_start_parser.add_argument( "--select-from-scope", action="store_true", help="AND-narrow the export selector with the persisted workspace scope", ) export_csv_start_parser.add_argument("--sort", "--sort-by", dest="sort", help="Sort field for search-based export or 'relevance'") export_csv_start_parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") export_csv_start_parser.add_argument("--budget-seconds", type=int, default=None, help="Cowork-safe status budget hint") export_csv_start_parser.add_argument( "--run-to-completion", action="store_true", help="Keep advancing the resumable CSV export until it reaches a terminal state", ) export_csv_run_step_parser = subparsers.add_parser( "export-csv-run-step", help="Run a bounded resumable CSV export step", ) export_csv_run_step_parser.add_argument("workspace", help="Workspace root path") export_csv_run_step_parser.add_argument("--run-id", required=True, help="Export run id") export_csv_run_step_parser.add_argument("--budget-seconds", type=int, default=None, help="Cowork-safe time budget") export_csv_status_parser = subparsers.add_parser("export-csv-status", help="Inspect a resumable CSV export") export_csv_status_parser.add_argument("workspace", help="Workspace root path") export_csv_status_parser.add_argument("--run-id", help="Export run id; defaults to the latest CSV export") export_csv_status_parser.add_argument("--budget-seconds", type=int, default=None, help="Cowork-safe status budget hint") export_archive_parser = subparsers.add_parser( "export-archive", help="Write selected documents, previews, and source artifacts to a zip archive", ) export_archive_parser.add_argument("workspace", help="Workspace root path") export_archive_parser.add_argument( "output_path", help="Zip file path; relative paths resolve under .retriever/exports", ) add_scope_run_selector_arguments(export_archive_parser) export_archive_parser.add_argument( "--family-mode", default="exact", choices=sorted(RUN_FAMILY_MODES), help="Whether to include only seed docs or their family members too", ) export_archive_parser.add_argument("--limit", dest="seed_limit", type=int, help="Limit the directly matched seed set") export_archive_parser.add_argument( "--portable-workspace", action="store_true", help="Include a curated subset .retriever/retriever.db for the exported documents", ) export_archive_start_parser = subparsers.add_parser( "export-archive-start", help="Start a resumable zip archive export", ) export_archive_start_parser.add_argument("workspace", help="Workspace root path") export_archive_start_parser.add_argument( "output_path", help="Zip file path; relative paths resolve under .retriever/exports", ) add_scope_run_selector_arguments(export_archive_start_parser) export_archive_start_parser.add_argument( "--family-mode", default="exact", choices=sorted(RUN_FAMILY_MODES), help="Whether to include only seed docs or their family members too", ) export_archive_start_parser.add_argument("--limit", dest="seed_limit", type=int, help="Limit the directly matched seed set") export_archive_start_parser.add_argument( "--portable-workspace", action="store_true", help="Include a curated subset .retriever/retriever.db for the exported documents", ) export_archive_start_parser.add_argument("--budget-seconds", type=int, default=None, help="Cowork-safe status budget hint") export_archive_start_parser.add_argument( "--run-to-completion", action="store_true", help="Keep advancing the resumable archive export until it reaches a terminal state", ) export_archive_run_step_parser = subparsers.add_parser( "export-archive-run-step", help="Run a bounded resumable zip archive export step", ) export_archive_run_step_parser.add_argument("workspace", help="Workspace root path") export_archive_run_step_parser.add_argument("--run-id", required=True, help="Export run id") export_archive_run_step_parser.add_argument("--budget-seconds", type=int, default=None, help="Cowork-safe time budget") export_archive_status_parser = subparsers.add_parser("export-archive-status", help="Inspect a resumable zip archive export") export_archive_status_parser.add_argument("workspace", help="Workspace root path") export_archive_status_parser.add_argument("--run-id", help="Export run id; defaults to the latest archive export") export_archive_status_parser.add_argument("--budget-seconds", type=int, default=None, help="Cowork-safe status budget hint") export_previews_parser = subparsers.add_parser( "export-previews", help="Write HTML preview exports for selected documents under .retriever/exports", ) export_previews_parser.add_argument("workspace", help="Workspace root path") export_previews_parser.add_argument("output_path", help="Output directory path; relative paths resolve under .retriever/exports") export_previews_parser.add_argument("query", nargs="?", default="", help="Optional keyword query text for search-based export") export_previews_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to export (repeatable, preserves input order)", ) export_previews_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable filter in the form ", ) export_previews_parser.add_argument("--sort", "--sort-by", dest="sort", help="Sort field for search-based export or 'relevance'") export_previews_parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") get_doc_parser = subparsers.add_parser("get-doc", help="Fetch one document with optional summary text or exact chunks") get_doc_parser.add_argument("workspace", help="Workspace root path") get_doc_parser.add_argument("--doc-id", dest="document_id", type=int, required=True, help="Document id") get_doc_parser.add_argument( "--include-text", choices=("none", "summary"), default="none", help="Include no extracted text or a deterministic summary prefix", ) get_doc_parser.add_argument( "--chunk", dest="chunk_indexes", action="append", type=int, help="Exact chunk index to include (repeatable)", ) get_doc_parser.add_argument( "--verbose", action="store_true", help="Return the full document context instead of the default compact JSON", ) list_chunks_parser = subparsers.add_parser("list-chunks", help="List chunk metadata for one document") list_chunks_parser.add_argument("workspace", help="Workspace root path") list_chunks_parser.add_argument("--doc-id", dest="document_id", type=int, required=True, help="Document id") list_chunks_parser.add_argument("--page", type=int, default=1, help="1-based chunk page") list_chunks_parser.add_argument( "--per-page", "--limit", dest="per_page", type=int, default=DEFAULT_CHUNK_PAGE_SIZE, help="Chunks per page", ) search_chunks_parser = subparsers.add_parser("search-chunks", help="Search matching text chunks with citations") search_chunks_parser.add_argument("workspace", help="Workspace root path") search_chunks_parser.add_argument("query", help="Keyword query text") search_chunks_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) search_chunks_parser.add_argument( "--sort", "--sort-by", dest="sort", choices=("relevance", "date_created", "date_modified"), help="Chunk result sort field", ) search_chunks_parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") search_chunks_parser.add_argument("--top-k", type=int, default=DEFAULT_CHUNK_SEARCH_TOP_K, help="Maximum chunks to return") search_chunks_parser.add_argument( "--per-doc-cap", type=int, default=DEFAULT_CHUNK_SEARCH_PER_DOC_CAP, help="Maximum chunks to keep from any one document", ) search_chunks_parser.add_argument("--count-only", action="store_true", help="Return counts instead of chunk payloads") search_chunks_parser.add_argument( "--distinct-docs", action="store_true", help="Count distinct documents with matching chunks when used with --count-only", ) search_chunks_parser.add_argument( "--select-from-scope", action="store_true", help="AND-narrow the chunk search to the persisted workspace scope", ) search_chunks_parser.add_argument( "--verbose", action="store_true", help="Return raw chunk text and the full document context instead of the default compact JSON", ) aggregate_parser = subparsers.add_parser("aggregate", help="Run bounded metadata aggregations across documents") aggregate_parser.add_argument("workspace", help="Workspace root path") aggregate_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) aggregate_parser.add_argument( "--group-by", dest="group_bys", action="append", required=True, help="Grouping expression, e.g. dataset_name or month:effective_date", ) aggregate_parser.add_argument("--metric", default="count", help="Aggregation metric") aggregate_parser.add_argument("--order-by", help="Bucket field name or 'metric'") aggregate_parser.add_argument("--order", choices=("asc", "desc"), help="Sort order") aggregate_parser.add_argument("--limit", type=int, default=DEFAULT_AGGREGATE_LIMIT, help="Maximum buckets to return") aggregate_parser.add_argument("--explain", action="store_true", help="Include generated SQL in the response") aggregate_parser.add_argument( "--select-from-scope", action="store_true", help="AND-narrow the aggregation to the persisted workspace scope", ) list_conversations_parser = subparsers.add_parser("list-conversations", help="List conversation summaries") list_conversations_parser.add_argument("workspace", help="Workspace root path") list_conversations_parser.add_argument("--query", help="Keyword query text") list_conversations_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) list_conversations_parser.add_argument( "--dataset", "--dataset-name", dest="dataset_names", action="append", help="Exact dataset name (repeatable)", ) list_conversations_parser.add_argument("--from-run-id", type=int, help="Restrict to documents already present in a prior run") list_conversations_parser.add_argument( "--select-from-scope", action="store_true", help="AND-narrow the conversation list to the persisted workspace scope", ) list_conversations_parser.add_argument("--limit", type=int, default=50, help="Maximum conversations to return") list_conversations_parser.add_argument("--offset", type=int, default=0, help="Zero-based conversation offset") list_conversations_parser.add_argument( "--sort", "--sort-by", dest="sort", help="Sort field: title, last_activity, first_activity, document_count, matching_document_count, source_kind, or conversation_type", ) list_conversations_parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") list_datasets_parser = subparsers.add_parser("list-datasets", help="List datasets in the workspace") list_datasets_parser.add_argument("workspace", help="Workspace root path") show_dataset_policy_parser = subparsers.add_parser("show-dataset-policy", help="Show source-backed entity merge policy for a dataset") show_dataset_policy_parser.add_argument("workspace", help="Workspace root path") add_dataset_selector_arguments(show_dataset_policy_parser) set_dataset_policy_parser = subparsers.add_parser("set-dataset-policy", help="Update source-backed entity merge policy for a dataset") set_dataset_policy_parser.add_argument("workspace", help="Workspace root path") add_dataset_selector_arguments(set_dataset_policy_parser) add_dataset_policy_boolean_argument( set_dataset_policy_parser, "--allow-auto-merge", dest="allow_auto_merge", help_text="Enable or disable all auto-merge resolution for this dataset", ) add_dataset_policy_boolean_argument( set_dataset_policy_parser, "--email-auto-merge", dest="email_auto_merge", help_text="Enable or disable exact email auto-merge", ) add_dataset_policy_boolean_argument( set_dataset_policy_parser, "--handle-auto-merge", dest="handle_auto_merge", help_text="Enable or disable provider-scoped handle auto-merge", ) add_dataset_policy_boolean_argument( set_dataset_policy_parser, "--phone-auto-merge", dest="phone_auto_merge", help_text="Enable or disable exact phone auto-merge", ) add_dataset_policy_boolean_argument( set_dataset_policy_parser, "--name-auto-merge", dest="name_auto_merge", help_text="Enable or disable exact normalized-name auto-merge", ) set_dataset_policy_parser.add_argument( "--external-id-auto-merge-name", dest="external_id_auto_merge_names", action="append", help="External identifier name eligible for exact auto-merge (repeatable; replaces the current list)", ) set_dataset_policy_parser.add_argument( "--clear-external-id-auto-merge-names", action="store_true", help="Clear all external-id auto-merge names", ) rebuild_entities_parser = subparsers.add_parser( "rebuild-entities", help="Rebuild entity recognition state from stored document metadata", ) rebuild_entities_parser.add_argument("workspace", help="Workspace root path") rebuild_entities_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Rebuild entity links for one document id (repeatable); omit for a full graph rebuild", ) rebuild_entities_parser.add_argument( "--batch-size", type=int, default=500, help="Documents to refresh per transaction", ) purge_vault_custodians_parser = subparsers.add_parser( "purge-vault-filename-custodians", help="Dry-run or apply cleanup for synthetic Google Vault MBOX filename custodians", ) purge_vault_custodians_parser.add_argument("workspace", help="Workspace root path") purge_vault_custodians_mode = purge_vault_custodians_parser.add_mutually_exclusive_group() purge_vault_custodians_mode.add_argument("--dry-run", action="store_true", help="List cleanup candidates without changing data") purge_vault_custodians_mode.add_argument("--apply", action="store_true", help="Apply the cleanup") entities_parser = subparsers.add_parser("entities", help="Entity maintenance commands") entities_subparsers = entities_parser.add_subparsers(dest="entity_command", required=True) entities_purge_vault_custodians_parser = entities_subparsers.add_parser( "purge-vault-filename-custodians", help="Dry-run or apply cleanup for synthetic Google Vault MBOX filename custodians", ) entities_purge_vault_custodians_parser.add_argument("workspace", help="Workspace root path") entities_purge_vault_custodians_mode = entities_purge_vault_custodians_parser.add_mutually_exclusive_group() entities_purge_vault_custodians_mode.add_argument("--dry-run", action="store_true", help="List cleanup candidates without changing data") entities_purge_vault_custodians_mode.add_argument("--apply", action="store_true", help="Apply the cleanup") rebuild_entities_start_parser = subparsers.add_parser( "rebuild-entities-start", help="Start a resumable entity rebuild run", ) rebuild_entities_start_parser.add_argument("workspace", help="Workspace root path") rebuild_entities_start_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Rebuild entity links for one document id (repeatable); omit for a full graph rebuild", ) rebuild_entities_start_parser.add_argument( "--batch-size", type=int, default=500, help="Documents to plan or refresh per bounded step", ) rebuild_entities_start_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Recommended per-step budget; values above the bounded-worker cap are rejected", ) rebuild_entities_status_parser = subparsers.add_parser( "rebuild-entities-status", help="Show resumable entity rebuild status", ) rebuild_entities_status_parser.add_argument("workspace", help="Workspace root path") rebuild_entities_status_parser.add_argument("--run-id", help="Run id; defaults to the latest entity rebuild run") rebuild_entities_status_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Budget used when computing next recommended commands", ) rebuild_entities_run_step_parser = subparsers.add_parser( "rebuild-entities-run-step", help="Advance a resumable entity rebuild until the bounded call budget is nearly exhausted", ) rebuild_entities_run_step_parser.add_argument("workspace", help="Workspace root path") rebuild_entities_run_step_parser.add_argument("--run-id", help="Run id to advance; defaults to the active or latest run") rebuild_entities_run_step_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Per-call budget; values above the bounded-worker cap are rejected", ) rebuild_entities_cancel_parser = subparsers.add_parser( "rebuild-entities-cancel", help="Cancel a resumable entity rebuild run", ) rebuild_entities_cancel_parser.add_argument("workspace", help="Workspace root path") rebuild_entities_cancel_parser.add_argument("--run-id", required=True, help="Run id to cancel") list_entities_parser = subparsers.add_parser("list-entities", help="List recognized entities") list_entities_parser.add_argument("workspace", help="Workspace root path") list_entities_parser.add_argument("--query", help="Filter entities by display name or identifier") list_entities_parser.add_argument("--limit", type=int, default=50, help="Maximum entities to return") list_entities_parser.add_argument("--offset", type=int, default=0, help="Zero-based entity offset") list_entities_parser.add_argument( "--sort", "--sort-by", dest="sort", help="Sort field: document_count, label, display_name, primary_email, entity_type, entity_origin, entity_status, or id", ) list_entities_parser.add_argument("--order", "--sort-order", dest="order", choices=("asc", "desc"), help="Sort order") list_entities_parser.add_argument( "--include-ignored", action="store_true", help="Include ignored entities; merged entities are still omitted", ) show_entity_parser = subparsers.add_parser("show-entity", help="Show one recognized entity") show_entity_parser.add_argument("workspace", help="Workspace root path") show_entity_parser.add_argument("entity_id", type=int, help="Entity id") show_entity_parser.add_argument("--limit", type=int, default=25, help="Maximum linked documents to return") create_entity_parser = subparsers.add_parser("create-entity", help="Create a manual entity") create_entity_parser.add_argument("workspace", help="Workspace root path") create_entity_parser.add_argument( "--entity-type", default=ENTITY_TYPE_PERSON, choices=sorted(ENTITY_TYPES), help="Entity type", ) create_entity_parser.add_argument("--display-name", help="Manual display name") create_entity_parser.add_argument("--notes", help="Optional entity notes") create_entity_parser.add_argument("--email", dest="emails", action="append", help="Email identifier (repeatable)") create_entity_parser.add_argument("--phone", dest="phones", action="append", help="Phone identifier (repeatable)") create_entity_parser.add_argument("--name", dest="names", action="append", help="Name identifier (repeatable)") create_entity_parser.add_argument( "--handle", dest="handles", action="append", help="Handle identifier as provider:scope:handle (repeatable)", ) create_entity_parser.add_argument( "--external-id", dest="external_ids", action="append", help="External identifier as name:value or name:scope:value (repeatable)", ) edit_entity_parser = subparsers.add_parser("edit-entity", help="Edit a manual entity profile") edit_entity_parser.add_argument("workspace", help="Workspace root path") edit_entity_parser.add_argument("entity_id", type=int, help="Entity id") edit_entity_parser.add_argument("--entity-type", choices=sorted(ENTITY_TYPES), help="Replacement entity type") edit_entity_display_group = edit_entity_parser.add_mutually_exclusive_group() edit_entity_display_group.add_argument("--display-name", help="Replacement manual display name") edit_entity_display_group.add_argument("--clear-display-name", action="store_true", help="Return display name to auto") edit_entity_notes_group = edit_entity_parser.add_mutually_exclusive_group() edit_entity_notes_group.add_argument("--notes", help="Replacement entity notes") edit_entity_notes_group.add_argument("--clear-notes", action="store_true", help="Clear entity notes") edit_entity_parser.add_argument("--add-email", dest="add_emails", action="append", help="Email identifier to add (repeatable)") edit_entity_parser.add_argument("--add-phone", dest="add_phones", action="append", help="Phone identifier to add (repeatable)") edit_entity_parser.add_argument("--add-name", dest="add_names", action="append", help="Name identifier to add (repeatable)") edit_entity_parser.add_argument( "--add-handle", dest="add_handles", action="append", help="Handle identifier as provider:scope:handle (repeatable)", ) edit_entity_parser.add_argument( "--add-external-id", dest="add_external_ids", action="append", help="External identifier as name:value or name:scope:value (repeatable)", ) entity_inventory_parser = subparsers.add_parser( "list-entity-role-inventory", help="List entity counts by role for a document scope", ) entity_inventory_parser.add_argument("workspace", help="Workspace root path") entity_inventory_parser.add_argument("query", nargs="?", default="", help="Optional keyword query text") entity_inventory_parser.add_argument( "--role", dest="roles", action="append", help="Entity role to include, e.g. author, participant, recipient, custodian (repeatable)", ) entity_inventory_parser.add_argument( "--filter", dest="filters", action="append", nargs="+", help="Repeatable SQL-like filter expression", ) entity_inventory_parser.add_argument( "--dataset", dest="dataset_names", action="append", help="Exact dataset name (repeatable)", ) entity_inventory_parser.add_argument("--dataset-id", type=int, help="Dataset id") entity_inventory_parser.add_argument("--conversation-id", type=int, help="Conversation id") entity_inventory_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to include (repeatable)", ) entity_inventory_parser.add_argument("--from-run-id", type=int, help="Restrict to documents from a prior run") entity_inventory_parser.add_argument("--limit", type=int, default=100, help="Maximum inventory rows to return") entity_inventory_parser.add_argument( "--examples", dest="examples_per_entity", type=int, default=3, help="Recent document examples per entity row", ) similar_entities_parser = subparsers.add_parser("similar-entities", help="Suggest active entities similar to one entity") similar_entities_parser.add_argument("workspace", help="Workspace root path") similar_entities_parser.add_argument("entity_id", type=int, help="Entity id") similar_entities_parser.add_argument("--limit", type=int, default=25, help="Maximum suggestions to return") merge_entities_parser = subparsers.add_parser("merge-entities", help="Merge one active entity into another") merge_entities_parser.add_argument("workspace", help="Workspace root path") merge_entities_parser.add_argument("source_entity_id", type=int, help="Entity id to merge away") merge_entities_parser.add_argument("target_entity_id", type=int, help="Entity id to keep") merge_entities_parser.add_argument("--force", action="store_true", help="Merge even when a merge block exists") merge_entities_parser.add_argument("--reason", help="Optional merge reason") block_entity_merge_parser = subparsers.add_parser("block-entity-merge", help="Prevent a suggested entity merge") block_entity_merge_parser.add_argument("workspace", help="Workspace root path") block_entity_merge_parser.add_argument("left_entity_id", type=int, help="First entity id") block_entity_merge_parser.add_argument("right_entity_id", type=int, help="Second entity id") block_entity_merge_parser.add_argument("--reason", help="Optional block reason") ignore_entity_parser = subparsers.add_parser("ignore-entity", help="Ignore a junk or non-entity record") ignore_entity_parser.add_argument("workspace", help="Workspace root path") ignore_entity_parser.add_argument("entity_id", type=int, help="Entity id") ignore_entity_parser.add_argument("--reason", help="Optional ignore reason") split_entity_parser = subparsers.add_parser("split-entity", help="Move selected identifiers or document links to another entity") split_entity_parser.add_argument("workspace", help="Workspace root path") split_entity_parser.add_argument("source_entity_id", type=int, help="Entity id to split from") split_entity_parser.add_argument("--target-entity-id", type=int, help="Existing target entity id; omitted creates a new entity") split_entity_parser.add_argument( "--identifier-id", dest="identifier_ids", action="append", type=int, help="Identifier id to move to the target (repeatable)", ) split_entity_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id whose source entity links should move to the target (repeatable)", ) split_entity_parser.add_argument( "--role", dest="roles", action="append", help="Role to move when --doc-id is used, e.g. author or recipient (repeatable)", ) split_entity_parser.add_argument("--display-name", help="Display name for a newly created split target") split_entity_parser.add_argument("--reason", help="Optional split reason") split_entity_parser.add_argument("--no-block", action="store_true", help="Do not create a merge block between source and target") assign_entity_parser = subparsers.add_parser("assign-entity", help="Manually assign an entity to a document role") assign_entity_parser.add_argument("workspace", help="Workspace root path") assign_entity_parser.add_argument("--doc-id", dest="document_id", type=int, required=True, help="Document id") assign_entity_parser.add_argument("--role", required=True, help="Entity role, e.g. author, participant, recipient, custodian") assign_entity_parser.add_argument("--entity-id", type=int, required=True, help="Entity id") assign_entity_parser.add_argument("--reason", help="Optional assignment reason") unassign_entity_parser = subparsers.add_parser("unassign-entity", help="Remove or suppress an entity link on a document role") unassign_entity_parser.add_argument("workspace", help="Workspace root path") unassign_entity_parser.add_argument("--doc-id", dest="document_id", type=int, required=True, help="Document id") unassign_entity_parser.add_argument("--role", required=True, help="Entity role, e.g. author, participant, recipient, custodian") unassign_entity_parser.add_argument("--entity-id", type=int, required=True, help="Entity id") unassign_entity_parser.add_argument("--reason", help="Optional unassignment reason") create_dataset_parser = subparsers.add_parser("create-dataset", help="Create a manual dataset") create_dataset_parser.add_argument("workspace", help="Workspace root path") create_dataset_parser.add_argument("dataset_name", help="Dataset name") add_to_dataset_parser = subparsers.add_parser("add-to-dataset", help="Add documents to a dataset") add_to_dataset_parser.add_argument("workspace", help="Workspace root path") add_to_dataset_parser.add_argument("--doc-id", dest="document_ids", action="append", type=int, required=True, help="Document id to add (repeatable)") add_dataset_selector_arguments(add_to_dataset_parser) remove_from_dataset_parser = subparsers.add_parser("remove-from-dataset", help="Remove documents from a dataset") remove_from_dataset_parser.add_argument("workspace", help="Workspace root path") remove_from_dataset_parser.add_argument("--doc-id", dest="document_ids", action="append", type=int, required=True, help="Document id to remove (repeatable)") add_dataset_selector_arguments(remove_from_dataset_parser) delete_dataset_parser = subparsers.add_parser("delete-dataset", help="Delete a dataset") delete_dataset_parser.add_argument("workspace", help="Workspace root path") add_dataset_selector_arguments(delete_dataset_parser) delete_docs_parser = subparsers.add_parser("delete-docs", help="Delete selected documents or matching occurrences") delete_docs_parser.add_argument("workspace", help="Workspace root path") delete_docs_parser.add_argument("--doc-id", dest="document_ids", action="append", type=int, help="Document id to delete (repeatable)") delete_docs_parser.add_argument( "--path", dest="path_prefixes", action="append", help="Delete documents whose rel_path matches this exact path or prefix (repeatable)", ) add_scope_run_selector_arguments(delete_docs_parser) delete_docs_parser.add_argument("--dry-run", action="store_true", help="Preview matches without deleting them") delete_docs_parser.add_argument("--confirm", action="store_true", help="Confirm the delete") list_runs_parser = subparsers.add_parser("list-runs", help="List planned processing runs") list_runs_parser.add_argument("workspace", help="Workspace root path") get_run_parser = subparsers.add_parser("get-run", help="Fetch one planned processing run") get_run_parser.add_argument("workspace", help="Workspace root path") get_run_parser.add_argument("--run-id", type=int, required=True, help="Run id") create_run_parser = subparsers.add_parser("create-run", help="Create a frozen processing run snapshot") create_run_parser.add_argument("workspace", help="Workspace root path") job_version_group = create_run_parser.add_mutually_exclusive_group(required=True) job_version_group.add_argument("--job-version-id", type=int, help="Explicit job version id") job_version_group.add_argument("--job-name", help="Job name") create_run_parser.add_argument("--job-version", dest="job_version_number", type=int, help="Job version number when selecting by job name") add_scope_run_selector_arguments(create_run_parser) create_run_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to include in the run (repeatable)", ) create_run_parser.add_argument( "--family-mode", default="exact", choices=sorted(RUN_FAMILY_MODES), help="Whether to include only seed docs or their family members too", ) create_run_parser.add_argument( "--activation-policy", default="manual", choices=sorted(RUN_ACTIVATION_POLICIES), help="Whether revision-producing jobs should auto-promote created text revisions", ) create_run_parser.add_argument("--limit", dest="seed_limit", type=int, help="Limit the directly matched seed set") list_text_revisions_parser = subparsers.add_parser("list-text-revisions", help="List stored text revisions for a document") list_text_revisions_parser.add_argument("workspace", help="Workspace root path") list_text_revisions_parser.add_argument("--doc-id", dest="document_id", type=int, required=True, help="Document id") activate_text_revision_parser = subparsers.add_parser("activate-text-revision", help="Promote a stored text revision to active indexed text") activate_text_revision_parser.add_argument("workspace", help="Workspace root path") activate_text_revision_parser.add_argument("--doc-id", dest="document_id", type=int, required=True, help="Document id") activate_text_revision_parser.add_argument("--text-revision-id", type=int, required=True, help="Stored text revision id") activate_text_revision_parser.add_argument( "--activation-policy", default="manual", choices=sorted(TEXT_REVISION_ACTIVATION_POLICIES), help="Audit label for why this revision is being promoted", ) list_results_parser = subparsers.add_parser("list-results", help="List stored processing results") list_results_parser.add_argument("workspace", help="Workspace root path") list_results_parser.add_argument("--run-id", type=int, help="Filter results to one run") list_results_parser.add_argument("--doc-id", dest="document_id", type=int, help="Filter results to one document") claim_run_items_parser = subparsers.add_parser("claim-run-items", help="Atomically claim pending run items for one worker") claim_run_items_parser.add_argument("workspace", help="Workspace root path") claim_run_items_parser.add_argument("--run-id", type=int, required=True, help="Run id") claim_run_items_parser.add_argument("--claimed-by", required=True, help="Worker/session identifier claiming the items") claim_run_items_parser.add_argument( "--limit", type=int, default=DEFAULT_RUN_ITEM_CLAIM_BATCH_SIZE, help="Maximum number of run items to claim", ) claim_run_items_parser.add_argument( "--stale-seconds", type=int, help=( "Reclaim running items whose heartbeat is older than this many seconds " f"(default: {DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS}s for inline, " f"{DEFAULT_RUN_ITEM_CLAIM_STALE_SECONDS}s for background)" ), ) claim_run_items_parser.add_argument( "--launch-mode", default="inline", choices=sorted(RUN_WORKER_MODES), help="Worker launch mode for supervision metadata", ) claim_run_items_parser.add_argument("--worker-task-id", help="Optional background task identifier") claim_run_items_parser.add_argument( "--max-batches", type=int, help="Optional maximum number of batches this worker should prepare before handing off", ) prepare_run_batch_parser = subparsers.add_parser( "prepare-run-batch", help="Claim one worker batch and return execution contexts plus worker hints", ) prepare_run_batch_parser.add_argument("workspace", help="Workspace root path") prepare_run_batch_parser.add_argument("--run-id", type=int, required=True, help="Run id") prepare_run_batch_parser.add_argument("--claimed-by", required=True, help="Worker/session identifier claiming the batch") prepare_run_batch_parser.add_argument( "--limit", type=int, help="Optional maximum number of run items to claim; defaults to the worker recommendation", ) prepare_run_batch_parser.add_argument( "--stale-seconds", type=int, help=( "Reclaim running items whose heartbeat is older than this many seconds " f"(default: {DEFAULT_COWORK_RUN_ITEM_CLAIM_STALE_SECONDS}s for inline, " f"{DEFAULT_RUN_ITEM_CLAIM_STALE_SECONDS}s for background)" ), ) prepare_run_batch_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Cowork-safe time budget for this batch preparation", ) prepare_run_batch_parser.add_argument( "--launch-mode", default="inline", choices=sorted(RUN_WORKER_MODES), help="Worker launch mode for supervision metadata", ) prepare_run_batch_parser.add_argument("--worker-task-id", help="Optional background task identifier") prepare_run_batch_parser.add_argument( "--max-batches", type=int, help="Optional maximum number of batches this worker should prepare before handing off", ) get_run_item_context_parser = subparsers.add_parser("get-run-item-context", help="Load the execution context for one run item") get_run_item_context_parser.add_argument("workspace", help="Workspace root path") get_run_item_context_parser.add_argument("--run-item-id", type=int, required=True, help="Run item id") heartbeat_run_items_parser = subparsers.add_parser("heartbeat-run-items", help="Refresh heartbeat timestamps for one worker's claimed items") heartbeat_run_items_parser.add_argument("workspace", help="Workspace root path") heartbeat_run_items_parser.add_argument("--run-id", type=int, required=True, help="Run id") heartbeat_run_items_parser.add_argument("--claimed-by", required=True, help="Worker/session identifier") finish_run_worker_parser = subparsers.add_parser( "finish-run-worker", help="Mark one worker as finished and persist its summary", ) finish_run_worker_parser.add_argument("workspace", help="Workspace root path") finish_run_worker_parser.add_argument("--run-id", type=int, required=True, help="Run id") finish_run_worker_parser.add_argument("--claimed-by", required=True, help="Worker/session identifier") finish_run_worker_parser.add_argument( "--worker-status", required=True, choices=sorted(status for status in RUN_WORKER_STATUSES if status != "active"), help="Terminal worker status", ) finish_run_worker_parser.add_argument("--summary-json", help="Optional worker summary JSON object") finish_run_worker_parser.add_argument("--error", dest="error_summary", help="Optional terminal error summary") complete_run_item_parser = subparsers.add_parser("complete-run-item", help="Mark one claimed run item completed and persist its result") complete_run_item_parser.add_argument("workspace", help="Workspace root path") complete_run_item_parser.add_argument("--run-item-id", type=int, required=True, help="Run item id") complete_run_item_parser.add_argument("--claimed-by", required=True, help="Worker/session identifier") complete_run_item_parser.add_argument( "--page-text", help="Plain-text completion for page-scoped visual run items (OCR or image description)", ) complete_run_item_parser.add_argument("--raw-output-json", help="Optional raw output JSON") complete_run_item_parser.add_argument("--normalized-output-json", help="Optional normalized output JSON") complete_run_item_parser.add_argument("--output-values-json", help="Optional per-output values JSON object") complete_run_item_parser.add_argument("--created-text-revision-json", help="Optional derived text revision payload JSON object") complete_run_item_parser.add_argument("--provider-metadata-json", help="Optional provider metadata JSON object") complete_run_item_parser.add_argument("--provider-request-id", help="Optional provider request identifier") complete_run_item_parser.add_argument("--input-tokens", type=int, help="Optional input token count") complete_run_item_parser.add_argument("--output-tokens", type=int, help="Optional output token count") complete_run_item_parser.add_argument("--cost-cents", type=int, help="Optional provider cost in cents") complete_run_item_parser.add_argument("--latency-ms", type=int, help="Optional execution latency in milliseconds") fail_run_item_parser = subparsers.add_parser("fail-run-item", help="Mark one claimed run item failed") fail_run_item_parser.add_argument("workspace", help="Workspace root path") fail_run_item_parser.add_argument("--run-item-id", type=int, required=True, help="Run item id") fail_run_item_parser.add_argument("--claimed-by", required=True, help="Worker/session identifier") fail_run_item_parser.add_argument("--error", required=True, help="Failure summary") fail_run_item_parser.add_argument("--provider-metadata-json", help="Optional provider metadata JSON object") fail_run_item_parser.add_argument("--provider-request-id", help="Optional provider request identifier") fail_run_item_parser.add_argument("--input-tokens", type=int, help="Optional input token count") fail_run_item_parser.add_argument("--output-tokens", type=int, help="Optional output token count") fail_run_item_parser.add_argument("--cost-cents", type=int, help="Optional provider cost in cents") fail_run_item_parser.add_argument("--latency-ms", type=int, help="Optional execution latency in milliseconds") run_status_parser = subparsers.add_parser("run-status", help="Summarize run progress, claims, and recent failures") run_status_parser.add_argument("workspace", help="Workspace root path") run_status_parser.add_argument("--run-id", type=int, required=True, help="Run id") run_status_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Budget used when rendering next_recommended_commands", ) run_job_step_parser = subparsers.add_parser( "run-job-step", help="Advance one Cowork-safe processing-run step or return one prepared worker batch", ) run_job_step_parser.add_argument("workspace", help="Workspace root path") run_job_step_parser.add_argument("--run-id", type=int, required=True, help="Run id") run_job_step_parser.add_argument("--claimed-by", help="Worker/session identifier; defaults to a stable Cowork id for the run") run_job_step_parser.add_argument( "--budget-seconds", type=int, default=DEFAULT_RESUMABLE_STEP_BUDGET_SECONDS, help="Cowork-safe time budget for this step", ) run_job_step_parser.add_argument("--limit", type=int, help="Optional maximum number of run items to claim") run_job_step_parser.add_argument( "--stale-seconds", type=int, help="Override stale claim recovery window for this step", ) run_job_step_parser.add_argument( "--launch-mode", default="inline", choices=sorted(RUN_WORKER_MODES), help="Worker launch mode for supervision metadata", ) run_job_step_parser.add_argument("--worker-task-id", help="Optional background task identifier") run_job_step_parser.add_argument( "--max-batches", type=int, help="Optional maximum number of batches this worker should prepare before handing off", ) cancel_run_parser = subparsers.add_parser("cancel-run", help="Stop claiming new work for a run and skip its pending items") cancel_run_parser.add_argument("workspace", help="Workspace root path") cancel_run_parser.add_argument("--run-id", type=int, required=True, help="Run id") cancel_run_parser.add_argument( "--force", action="store_true", help="Request force-stop for background workers that expose task ids", ) finalize_ocr_run_parser = subparsers.add_parser("finalize-ocr-run", help="Merge completed OCR page items into document-level OCR results") finalize_ocr_run_parser.add_argument("workspace", help="Workspace root path") finalize_ocr_run_parser.add_argument("--run-id", type=int, required=True, help="Run id") finalize_ocr_run_parser.add_argument( "--cascade-metadata", action="store_true", help="Immediately refresh text-derived metadata for docs whose active OCR text changed", ) finalize_ocr_run_parser.add_argument( "--metadata-field", dest="metadata_fields", action="append", choices=sorted(TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN), help="Metadata field to refresh when --cascade-metadata is used", ) finalize_image_description_run_parser = subparsers.add_parser( "finalize-image-description-run", help="Merge completed image-description page items into document-level text revisions", ) finalize_image_description_run_parser.add_argument("workspace", help="Workspace root path") finalize_image_description_run_parser.add_argument("--run-id", type=int, required=True, help="Run id") publish_run_results_parser = subparsers.add_parser( "publish-run-results", help="Publish bound result outputs from a run into custom fields", ) publish_run_results_parser.add_argument("workspace", help="Workspace root path") publish_run_results_parser.add_argument("--run-id", type=int, required=True, help="Run id") publish_run_results_parser.add_argument( "--output-name", dest="output_names", action="append", help="Optional job output name to publish (repeatable)", ) list_jobs_parser = subparsers.add_parser("list-jobs", help="List configured processing jobs") list_jobs_parser.add_argument("workspace", help="Workspace root path") create_job_parser = subparsers.add_parser("create-job", help="Create a processing job") create_job_parser.add_argument("workspace", help="Workspace root path") create_job_parser.add_argument("job_name", help="Job name") create_job_parser.add_argument("job_kind", choices=sorted(JOB_KINDS), help="Job kind") create_job_parser.add_argument("--description", help="Optional job description") add_job_output_parser = subparsers.add_parser("add-job-output", help="Create or update a job output") add_job_output_parser.add_argument("workspace", help="Workspace root path") add_job_output_parser.add_argument("job_name", help="Existing job name") add_job_output_parser.add_argument("output_name", help="Job output name") add_job_output_parser.add_argument( "--value-type", default="text", choices=sorted(JOB_OUTPUT_VALUE_TYPES), help="Logical output value type", ) add_job_output_parser.add_argument("--bind-custom-field", dest="bound_custom_field", help="Optional custom field binding") add_job_output_parser.add_argument("--description", help="Optional output description") list_job_versions_parser = subparsers.add_parser("list-job-versions", help="List versions for one job") list_job_versions_parser.add_argument("workspace", help="Workspace root path") list_job_versions_parser.add_argument("job_name", help="Existing job name") create_job_version_parser = subparsers.add_parser("create-job-version", help="Create a new immutable job version") create_job_version_parser.add_argument("workspace", help="Workspace root path") create_job_version_parser.add_argument("job_name", help="Existing job name") create_job_version_parser.add_argument("--instruction", help="Optional job instruction text") create_job_version_parser.add_argument( "--capability", choices=sorted(JOB_CAPABILITIES), help="Optional Cowork execution capability; defaults from job kind when omitted", ) create_job_version_parser.add_argument( "--provider", default="cowork_agent", help="Optional provider identifier (defaults to cowork_agent; external providers are future-facing)", ) create_job_version_parser.add_argument("--model", help="Optional model name") create_job_version_parser.add_argument( "--input-basis", choices=sorted(JOB_INPUT_BASES), help="Primary input basis for this version; defaults from job kind when omitted", ) create_job_version_parser.add_argument("--response-schema-json", help="Optional JSON schema payload") create_job_version_parser.add_argument("--parameters-json", help="Optional provider parameters as JSON object") create_job_version_parser.add_argument("--segment-profile", help="Optional segment profile name") create_job_version_parser.add_argument("--aggregation-strategy", help="Optional aggregation strategy") create_job_version_parser.add_argument("--display-name", help="Optional display name override") list_fields_parser = subparsers.add_parser("list-fields", help="List registered custom document fields") list_fields_parser.add_argument("workspace", help="Workspace root path") list_fields_parser.add_argument( "--format", choices=("json", "table"), default="json", help="Output format", ) add_field_parser = subparsers.add_parser("add-field", help="Add a custom document field") add_field_parser.add_argument("workspace", help="Workspace root path") add_field_parser.add_argument("field_name", help="Field name") add_field_parser.add_argument("field_type", choices=sorted(REGISTRY_FIELD_TYPES), help="Field type") add_field_parser.add_argument("--instruction", help="Field extraction instruction") rename_field_parser = subparsers.add_parser("rename-field", help="Rename an existing custom field") rename_field_parser.add_argument("workspace", help="Workspace root path") rename_field_parser.add_argument("old_name", help="Existing custom field name") rename_field_parser.add_argument("new_name", help="New custom field name") delete_field_parser = subparsers.add_parser("delete-field", help="Delete an existing custom field") delete_field_parser.add_argument("workspace", help="Workspace root path") delete_field_parser.add_argument("field_name", help="Existing custom field name") delete_field_parser.add_argument("--confirm", action="store_true", help="Confirm the irreversible delete") describe_field_parser = subparsers.add_parser("describe-field", help="Set or clear a custom field description") describe_field_parser.add_argument("workspace", help="Workspace root path") describe_field_parser.add_argument("field_name", help="Existing custom field name") describe_group = describe_field_parser.add_mutually_exclusive_group(required=True) describe_group.add_argument("--text", help="Replacement description text") describe_group.add_argument("--clear", action="store_true", help="Clear the existing description") change_field_type_parser = subparsers.add_parser("change-field-type", help="Change a custom field type in place") change_field_type_parser.add_argument("workspace", help="Workspace root path") change_field_type_parser.add_argument("field_name", help="Existing custom field name") change_field_type_parser.add_argument( "target_field_type", choices=sorted(REGISTRY_FIELD_TYPES), help="Target field type", ) promote_field_parser = subparsers.add_parser("promote-field-type", help=argparse.SUPPRESS) promote_field_parser.add_argument("workspace", help="Workspace root path") promote_field_parser.add_argument("field_name", help="Existing custom field name") promote_field_parser.add_argument("target_field_type", choices=("date",), help="Target field type") fill_field_parser = subparsers.add_parser("fill-field", help="Set or clear a field value on one or more documents") fill_field_parser.add_argument("workspace", help="Workspace root path") fill_field_parser.add_argument("--field", required=True, help="Field name") fill_value_group = fill_field_parser.add_mutually_exclusive_group(required=True) fill_value_group.add_argument("--value", help="Replacement field value") fill_value_group.add_argument("--clear", action="store_true", help="Clear the field value") add_scope_run_selector_arguments(fill_field_parser) fill_field_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to update (repeatable)", ) fill_field_parser.add_argument("--dry-run", action="store_true", help="Preview matching documents without writing") fill_field_parser.add_argument("--confirm", action="store_true", help="Confirm bulk writes") set_field_parser = subparsers.add_parser("set-field", help=argparse.SUPPRESS) set_field_parser.add_argument("workspace", help="Workspace root path") set_field_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, required=True, help="Document id to update (repeatable)", ) set_field_parser.add_argument("--field", required=True, help="Field name") set_field_parser.add_argument("--value", help="Field value") refresh_text_metadata_parser = subparsers.add_parser( "refresh-text-derived-metadata", help="Re-derive author/participants/date-created metadata from active document text", ) refresh_text_metadata_parser.add_argument("workspace", help="Workspace root path") add_scope_run_selector_arguments(refresh_text_metadata_parser) refresh_text_metadata_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to refresh (repeatable)", ) refresh_text_metadata_parser.add_argument( "--field", dest="field_names", action="append", choices=sorted(TEXT_DERIVED_METADATA_FIELD_TO_OCCURRENCE_COLUMN), help="Metadata field to refresh (repeatable; defaults to author, participants, date_created)", ) reconcile_duplicates_parser = subparsers.add_parser( "reconcile-duplicates", help="Dry-run or apply post-ingest duplicate reconciliation", ) reconcile_duplicates_parser.add_argument("workspace", help="Workspace root path") reconcile_duplicates_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to consider for reconciliation (repeatable)", ) add_scope_run_selector_arguments(reconcile_duplicates_parser) reconcile_duplicates_parser.add_argument( "--basis", choices=("content_hash",), default="content_hash", help="Reconciliation basis", ) reconcile_mode_group = reconcile_duplicates_parser.add_mutually_exclusive_group() reconcile_mode_group.add_argument("--dry-run", action="store_true", help="Preview merge candidates without applying them") reconcile_mode_group.add_argument("--apply", action="store_true", help="Apply mergeable reconciliation candidates") merge_conversation_parser = subparsers.add_parser( "merge-into-conversation", help="Manually assign one document family into another document's conversation", ) merge_conversation_parser.add_argument("workspace", help="Workspace root path") merge_conversation_parser.add_argument("--doc-id", type=int, required=True, help="Document id to move") merge_conversation_parser.add_argument("--target-doc-id", type=int, required=True, help="Target document id") split_conversation_parser = subparsers.add_parser( "split-from-conversation", help="Split one document family into its own manually pinned conversation", ) split_conversation_parser.add_argument("workspace", help="Workspace root path") split_conversation_parser.add_argument("--doc-id", type=int, required=True, help="Document id to split") clear_conversation_parser = subparsers.add_parser( "clear-conversation-assignment", help="Clear a manual conversation pin and re-run automatic conversation assignment", ) clear_conversation_parser.add_argument("workspace", help="Workspace root path") clear_conversation_parser.add_argument("--doc-id", type=int, required=True, help="Document id to clear") refresh_previews_parser = subparsers.add_parser( "refresh-previews", aliases=["refresh-conversation-previews"], help="Regenerate generated document and conversation preview artifacts", ) refresh_previews_parser.add_argument("workspace", help="Workspace root path") refresh_previews_parser.add_argument( "--scope", choices=["conversations", "documents", "all"], default="conversations", help="Preview family to refresh", ) refresh_previews_parser.add_argument( "--conversation-id", dest="conversation_ids", action="append", type=int, help="Conversation id to refresh (repeatable; conversation scope)", ) refresh_previews_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id to refresh or use as a conversation selector (repeatable)", ) refresh_dataset_selector_group = refresh_previews_parser.add_mutually_exclusive_group() refresh_dataset_selector_group.add_argument("--dataset-id", type=int, help="Dataset id") refresh_dataset_selector_group.add_argument("--dataset-name", help="Exact dataset name") refresh_previews_parser.add_argument( "--missing-only", action="store_true", help="Only refresh preview packages with missing rows or files", ) refresh_previews_parser.add_argument( "--from-source", action="store_true", help="Request source-backed preview regeneration where the selected scope supports it", ) rebuild_conversations_parser = subparsers.add_parser( "rebuild-conversations", help="Re-run conversation assignment and regenerate conversation previews", ) rebuild_conversations_parser.add_argument("workspace", help="Workspace root path") rebuild_conversations_parser.add_argument( "--conversation-id", dest="conversation_ids", action="append", type=int, help="Conversation id to refresh after reassignment (repeatable)", ) rebuild_conversations_parser.add_argument( "--doc-id", dest="document_ids", action="append", type=int, help="Document id whose conversation should be refreshed after reassignment (repeatable)", ) rebuild_conversations_dataset_group = rebuild_conversations_parser.add_mutually_exclusive_group() rebuild_conversations_dataset_group.add_argument("--dataset-id", type=int, help="Dataset id") rebuild_conversations_dataset_group.add_argument("--dataset-name", help="Exact dataset name") rebuild_conversations_parser.add_argument( "--batch-size", type=int, default=50, help="Conversations to refresh per transaction", ) subparsers.add_parser("schema-version", help="Print the schema version") return parser def _auto_upgrade_and_maybe_reexec(root: Path, command: str) -> None: """Best-effort runtime metadata refresh for older workspaces.""" if command in AUTO_UPGRADE_EXEMPT_COMMANDS: return result = maybe_upgrade_workspace_tool(root) if not result: return try: print( "retriever-runtime-sync: " + json.dumps(result, sort_keys=True), file=sys.stderr, ) except Exception: # pragma: no cover - never fail dispatch over logging pass def main() -> int: benchmark_mark("main_entered") parser = build_parser() args = parser.parse_args() if bool(getattr(args, "human", False)): args.output = "human" set_cli_output_mode(str(getattr(args, "output", "json"))) benchmark_mark("argparse_done", command=getattr(args, "command", None)) try: if args.command == "schema-version": print(json.dumps({"schema_version": SCHEMA_VERSION, "tool_version": TOOL_VERSION})) return 0 root = Path(args.workspace).expanduser().resolve() set_active_workspace_root(root) if args.command == "workspace": if args.workspace_action == "status": status_payload = workspace_status(root, bool(getattr(args, "quick", False))) return emit_cli_payload("workspace", {"action": "status", **status_payload}) if args.workspace_action == "init": return emit_cli_payload( "workspace", init_workspace(root, quick=bool(getattr(args, "quick", False))), ) if args.workspace_action == "update": canonical_source = getattr(args, "canonical_source", None) if canonical_source: canonical_path = Path(canonical_source).expanduser().resolve() if not canonical_path.is_file(): raise RetrieverError( f"Canonical tool not found at --from path: {canonical_path}" ) else: canonical_path = locate_canonical_plugin_tool_or_self() if canonical_path is None: raise RetrieverError( "Could not auto-discover the canonical tools.py bundle. " "Pass --from or set RETRIEVER_CANONICAL_TOOL_PATH." ) update_payload = upgrade_workspace_tool( root, canonical_path, force=bool(getattr(args, "force", False)), reason="manual", ) return emit_cli_payload("workspace", {"action": "update", **update_payload}) if args.workspace_action == "doctor": doctor_payload = doctor( root, bool(getattr(args, "quick", False)), repair_stale_sidecars=bool(getattr(args, "repair_stale_sidecars", False)), ) return emit_cli_payload("workspace", {"action": "doctor", **doctor_payload}) _auto_upgrade_and_maybe_reexec(root, args.command) if args.command == "ingest": ingest_progress = LongTaskProgressPrinter("ingest") if bool(args.run_to_completion) else None return emit_cli_payload( "ingest", ingest_v2_facade( root, recursive=args.recursive, raw_file_types=args.file_types, raw_paths=args.paths, budget_seconds=args.budget_seconds, run_to_completion=args.run_to_completion, progress_callback=ingest_progress.emit if ingest_progress is not None else None, ), ) if args.command == "ingest-start": return emit_cli_payload( "ingest-start", ingest_v2_start( root, recursive=args.recursive, raw_file_types=args.file_types, raw_paths=args.paths, budget_seconds=args.budget_seconds, ), ) if args.command == "ingest-status": return emit_cli_payload( "ingest-status", ingest_v2_status(root, run_id=args.run_id, budget_seconds=args.budget_seconds), ) if args.command == "ingest-cancel": return emit_cli_payload( "ingest-cancel", ingest_v2_cancel(root, run_id=args.run_id, force=args.force), ) if args.command == "ingest-run-step": return emit_cli_payload( "ingest-run-step", ingest_v2_run_step(root, run_id=args.run_id, budget_seconds=args.budget_seconds), ) if args.command == "ingest-plan-step": return emit_cli_payload( "ingest-plan-step", ingest_v2_plan_step(root, run_id=args.run_id, budget_seconds=args.budget_seconds), ) if args.command == "ingest-prepare-step": return emit_cli_payload( "ingest-prepare-step", ingest_v2_prepare_step(root, run_id=args.run_id, budget_seconds=args.budget_seconds), ) if args.command == "ingest-commit-step": return emit_cli_payload( "ingest-commit-step", ingest_v2_commit_step( root, run_id=args.run_id, budget_seconds=args.budget_seconds, max_items=args.max_items, ), ) if args.command == "ingest-finalize-step": return emit_cli_payload( "ingest-finalize-step", ingest_v2_finalize_step(root, run_id=args.run_id, budget_seconds=args.budget_seconds), ) if args.command == "ingest-production": return emit_cli_payload("ingest-production", ingest_production(root, args.production_root)) if args.command == "inspect-pst-properties": print( json.dumps( inspect_pst_properties( root, args.pst_path, source_item_ids=args.source_item_ids, message_kind=args.message_kind, limit=args.limit, max_record_entries=args.max_record_entries, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "search": return emit_cli_payload( "search", search( root, args.query, args.filters, args.sort, args.order, args.page, args.per_page, args.columns, args.mode, compact_mode=(not args.verbose and args.mode == "compose"), ), verbose=args.verbose, ) if args.command == "slash": raw_command = " ".join(args.command_text).strip() payload = run_slash_command(root, raw_command) rendered_output = render_slash_read_only_output(raw_command, payload) if rendered_output is not None: sys.stdout.write(rendered_output + "\n") sys.stdout.flush() else: print(json.dumps(payload, indent=2, sort_keys=True)) return 0 if args.command == "catalog": return emit_cli_payload("catalog", catalog(root)) if args.command == "export-csv": export_csv_progress = LongTaskProgressPrinter("export-csv") print( json.dumps( export_csv( root, args.output_path, args.fields, args.document_ids, args.query, args.filters, args.sort, args.order, args.select_from_scope, progress_callback=export_csv_progress.emit, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "export-csv-start": export_csv_progress = LongTaskProgressPrinter("export-csv-start") if bool(args.run_to_completion) else None return emit_cli_payload( "export-csv-start", export_csv_start( root, args.output_path, args.fields, args.document_ids, args.query, args.filters, args.sort, args.order, args.select_from_scope, budget_seconds=args.budget_seconds, run_to_completion=args.run_to_completion, progress_callback=export_csv_progress.emit if export_csv_progress is not None else None, ), ) if args.command == "export-csv-run-step": return emit_cli_payload( "export-csv-run-step", export_run_step( root, export_kind="csv", run_id=args.run_id, budget_seconds=args.budget_seconds, ), ) if args.command == "export-csv-status": return emit_cli_payload( "export-csv-status", export_status( root, export_kind="csv", run_id=args.run_id, budget_seconds=args.budget_seconds, ), ) if args.command == "export-archive": export_archive_progress = LongTaskProgressPrinter("export-archive") print( json.dumps( export_archive( root, args.output_path, dataset_names=args.dataset_names, query=args.query, raw_bates=args.bates, raw_filters=args.filters, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, family_mode=args.family_mode, seed_limit=args.seed_limit, portable_workspace=args.portable_workspace, progress_callback=export_archive_progress.emit, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "export-archive-start": export_archive_progress = LongTaskProgressPrinter("export-archive-start") if bool(args.run_to_completion) else None return emit_cli_payload( "export-archive-start", export_archive_start( root, args.output_path, dataset_names=args.dataset_names, query=args.query, raw_bates=args.bates, raw_filters=args.filters, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, family_mode=args.family_mode, seed_limit=args.seed_limit, portable_workspace=args.portable_workspace, budget_seconds=args.budget_seconds, run_to_completion=args.run_to_completion, progress_callback=export_archive_progress.emit if export_archive_progress is not None else None, ), ) if args.command == "export-archive-run-step": return emit_cli_payload( "export-archive-run-step", export_run_step( root, export_kind="archive", run_id=args.run_id, budget_seconds=args.budget_seconds, ), ) if args.command == "export-archive-status": return emit_cli_payload( "export-archive-status", export_status( root, export_kind="archive", run_id=args.run_id, budget_seconds=args.budget_seconds, ), ) if args.command == "export-previews": print( json.dumps( export_previews( root, args.output_path, args.document_ids, args.query, args.filters, args.sort, args.order, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "get-doc": return emit_cli_payload( "get-doc", get_doc(root, args.document_id, args.include_text, args.chunk_indexes), verbose=args.verbose, ) if args.command == "list-chunks": return emit_cli_payload("list-chunks", list_chunks(root, args.document_id, args.page, args.per_page)) if args.command == "search-chunks": return emit_cli_payload( "search-chunks", search_chunks( root, args.query, args.filters, args.sort, args.order, args.top_k, args.per_doc_cap, select_from_scope=args.select_from_scope, count_only=args.count_only, distinct_docs=args.distinct_docs, ), verbose=args.verbose, ) if args.command == "aggregate": return emit_cli_payload( "aggregate", aggregate( root, args.filters, args.group_bys, args.metric, args.order_by, args.order, args.limit, args.explain, select_from_scope=args.select_from_scope, ), ) if args.command == "list-conversations": return emit_cli_payload( "list-conversations", list_conversations( root, query=args.query, raw_filters=args.filters, dataset_names=args.dataset_names, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, limit=args.limit, offset=args.offset, sort=args.sort, order=args.order, ), ) if args.command == "list-datasets": return emit_cli_payload("list-datasets", list_datasets(root)) if args.command == "show-dataset-policy": return emit_cli_payload( "show-dataset-policy", show_dataset_policy( root, dataset_id=args.dataset_id, dataset_name=args.dataset_name, ), ) if args.command == "set-dataset-policy": return emit_cli_payload( "set-dataset-policy", set_dataset_policy( root, dataset_id=args.dataset_id, dataset_name=args.dataset_name, allow_auto_merge=args.allow_auto_merge, email_auto_merge=args.email_auto_merge, handle_auto_merge=args.handle_auto_merge, phone_auto_merge=args.phone_auto_merge, name_auto_merge=args.name_auto_merge, external_id_auto_merge_names=args.external_id_auto_merge_names, clear_external_id_auto_merge_names=args.clear_external_id_auto_merge_names, ), ) if args.command == "rebuild-entities": rebuild_entities_progress = LongTaskProgressPrinter("rebuild-entities") return emit_cli_payload( "rebuild-entities", rebuild_entities( root, document_ids=args.document_ids, batch_size=args.batch_size, progress_callback=rebuild_entities_progress.emit, ), ) if args.command == "purge-vault-filename-custodians": return emit_cli_payload( "purge-vault-filename-custodians", purge_vault_filename_custodians(root, apply=bool(args.apply)), ) if args.command == "entities": if args.entity_command == "purge-vault-filename-custodians": return emit_cli_payload( "entities purge-vault-filename-custodians", purge_vault_filename_custodians(root, apply=bool(args.apply)), ) if args.command == "rebuild-entities-start": return emit_cli_payload( "rebuild-entities-start", rebuild_entities_start( root, document_ids=args.document_ids, batch_size=args.batch_size, budget_seconds=args.budget_seconds, ), ) if args.command == "rebuild-entities-status": return emit_cli_payload( "rebuild-entities-status", rebuild_entities_status( root, run_id=args.run_id, budget_seconds=args.budget_seconds, ), ) if args.command == "rebuild-entities-run-step": return emit_cli_payload( "rebuild-entities-run-step", rebuild_entities_run_step( root, run_id=args.run_id, budget_seconds=args.budget_seconds, ), ) if args.command == "rebuild-entities-cancel": return emit_cli_payload( "rebuild-entities-cancel", rebuild_entities_cancel( root, run_id=args.run_id, ), ) if args.command == "list-entities": return emit_cli_payload( "list-entities", run_entity_list_command( root, query=args.query, limit=args.limit, offset=args.offset, sort=args.sort, order=args.order, include_ignored=args.include_ignored, ), ) if args.command == "show-entity": return emit_cli_payload( "show-entity", show_entity(root, args.entity_id, document_limit=args.limit), ) if args.command == "create-entity": return emit_cli_payload( "create-entity", create_entity( root, entity_type=args.entity_type, display_name=args.display_name, notes=args.notes, emails=args.emails, phones=args.phones, names=args.names, handles=args.handles, external_ids=args.external_ids, ), ) if args.command == "edit-entity": return emit_cli_payload( "edit-entity", edit_entity( root, args.entity_id, entity_type=args.entity_type, display_name=args.display_name, clear_display_name=args.clear_display_name, notes=args.notes, clear_notes=args.clear_notes, add_emails=args.add_emails, add_phones=args.add_phones, add_names=args.add_names, add_handles=args.add_handles, add_external_ids=args.add_external_ids, ), ) if args.command == "list-entity-role-inventory": return emit_cli_payload( "list-entity-role-inventory", list_entity_role_inventory( root, roles=args.roles, query=args.query, raw_filters=args.filters, document_ids=args.document_ids, dataset_id=args.dataset_id, dataset_names=args.dataset_names, conversation_id=args.conversation_id, from_run_id=args.from_run_id, limit=args.limit, examples_per_entity=args.examples_per_entity, ), ) if args.command == "similar-entities": return emit_cli_payload( "similar-entities", similar_entities(root, args.entity_id, limit=args.limit), ) if args.command == "merge-entities": return emit_cli_payload( "merge-entities", merge_entities( root, args.source_entity_id, args.target_entity_id, force=args.force, reason=args.reason, ), ) if args.command == "block-entity-merge": return emit_cli_payload( "block-entity-merge", block_entity_merge( root, args.left_entity_id, args.right_entity_id, reason=args.reason, ), ) if args.command == "ignore-entity": return emit_cli_payload( "ignore-entity", ignore_entity(root, args.entity_id, reason=args.reason), ) if args.command == "split-entity": return emit_cli_payload( "split-entity", split_entity( root, args.source_entity_id, target_entity_id=args.target_entity_id, identifier_ids=args.identifier_ids, document_ids=args.document_ids, roles=args.roles, display_name=args.display_name, reason=args.reason, block_merge=not args.no_block, ), ) if args.command == "assign-entity": return emit_cli_payload( "assign-entity", assign_entity( root, document_id=args.document_id, role=args.role, entity_id=args.entity_id, reason=args.reason, ), ) if args.command == "unassign-entity": return emit_cli_payload( "unassign-entity", unassign_entity( root, document_id=args.document_id, role=args.role, entity_id=args.entity_id, reason=args.reason, ), ) if args.command == "create-dataset": return emit_cli_payload("create-dataset", create_dataset(root, args.dataset_name)) if args.command == "add-to-dataset": return emit_cli_payload( "add-to-dataset", add_to_dataset( root, args.document_ids, dataset_id=args.dataset_id, dataset_name=args.dataset_name, ), ) if args.command == "remove-from-dataset": return emit_cli_payload( "remove-from-dataset", remove_from_dataset( root, args.document_ids, dataset_id=args.dataset_id, dataset_name=args.dataset_name, ), ) if args.command == "delete-dataset": return emit_cli_payload( "delete-dataset", delete_dataset( root, dataset_id=args.dataset_id, dataset_name=args.dataset_name, ), ) if args.command == "delete-docs": return emit_cli_payload( "delete-docs", delete_docs( root, document_ids=args.document_ids, query=args.query, raw_bates=args.bates, raw_filters=args.filters, dataset_names=args.dataset_names, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, path_prefixes=args.path_prefixes, dry_run=args.dry_run, confirm=args.confirm, ), ) if args.command == "list-runs": print(json.dumps(list_runs(root), indent=2, sort_keys=True)) return 0 if args.command == "get-run": print(json.dumps(get_run(root, args.run_id), indent=2, sort_keys=True)) return 0 if args.command == "create-run": print( json.dumps( create_run( root, job_version_id=args.job_version_id, raw_job_name=args.job_name, job_version_number=args.job_version_number, dataset_names=args.dataset_names, document_ids=args.document_ids, query=args.query, raw_bates=args.bates, raw_filters=args.filters, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, activation_policy=args.activation_policy, family_mode=args.family_mode, seed_limit=args.seed_limit, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "list-text-revisions": print(json.dumps(list_text_revisions(root, document_id=args.document_id), indent=2, sort_keys=True)) return 0 if args.command == "activate-text-revision": print( json.dumps( activate_text_revision( root, document_id=args.document_id, text_revision_id=args.text_revision_id, activation_policy=args.activation_policy, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "list-results": print( json.dumps( list_results(root, run_id=args.run_id, document_id=args.document_id), indent=2, sort_keys=True, ) ) return 0 if args.command == "claim-run-items": print( json.dumps( claim_run_items( root, run_id=args.run_id, claimed_by=args.claimed_by, limit=args.limit, stale_after_seconds=args.stale_seconds, launch_mode=args.launch_mode, worker_task_id=args.worker_task_id, max_batches=args.max_batches, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "prepare-run-batch": print( json.dumps( prepare_run_batch( root, run_id=args.run_id, claimed_by=args.claimed_by, limit=args.limit, stale_after_seconds=args.stale_seconds, launch_mode=args.launch_mode, worker_task_id=args.worker_task_id, max_batches=args.max_batches, budget_seconds=args.budget_seconds, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "get-run-item-context": print(json.dumps(get_run_item_context(root, run_item_id=args.run_item_id), indent=2, sort_keys=True)) return 0 if args.command == "heartbeat-run-items": print( json.dumps( heartbeat_run_items(root, run_id=args.run_id, claimed_by=args.claimed_by), indent=2, sort_keys=True, ) ) return 0 if args.command == "finish-run-worker": print( json.dumps( finish_run_worker( root, run_id=args.run_id, claimed_by=args.claimed_by, worker_status=args.worker_status, summary_json=args.summary_json, error_summary=args.error_summary, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "complete-run-item": print( json.dumps( complete_run_item( root, run_item_id=args.run_item_id, claimed_by=args.claimed_by, page_text=args.page_text, raw_output_json=args.raw_output_json, normalized_output_json=args.normalized_output_json, output_values_json=args.output_values_json, created_text_revision_json=args.created_text_revision_json, provider_metadata_json=args.provider_metadata_json, provider_request_id=args.provider_request_id, input_tokens=args.input_tokens, output_tokens=args.output_tokens, cost_cents=args.cost_cents, latency_ms=args.latency_ms, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "fail-run-item": print( json.dumps( fail_run_item( root, run_item_id=args.run_item_id, claimed_by=args.claimed_by, error_summary=args.error, provider_metadata_json=args.provider_metadata_json, provider_request_id=args.provider_request_id, input_tokens=args.input_tokens, output_tokens=args.output_tokens, cost_cents=args.cost_cents, latency_ms=args.latency_ms, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "run-status": print( json.dumps( run_status(root, run_id=args.run_id, budget_seconds=args.budget_seconds), indent=2, sort_keys=True, ) ) return 0 if args.command == "run-job-step": print( json.dumps( run_job_step( root, run_id=args.run_id, claimed_by=args.claimed_by, budget_seconds=args.budget_seconds, limit=args.limit, stale_after_seconds=args.stale_seconds, launch_mode=args.launch_mode, worker_task_id=args.worker_task_id, max_batches=args.max_batches, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "cancel-run": print(json.dumps(cancel_run(root, run_id=args.run_id, force=args.force), indent=2, sort_keys=True)) return 0 if args.command == "finalize-ocr-run": print( json.dumps( finalize_ocr_run( root, run_id=args.run_id, cascade_metadata=bool(args.cascade_metadata), metadata_fields=args.metadata_fields, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "finalize-image-description-run": print(json.dumps(finalize_image_description_run(root, run_id=args.run_id), indent=2, sort_keys=True)) return 0 if args.command == "publish-run-results": print( json.dumps( publish_run_results(root, run_id=args.run_id, raw_output_names=args.output_names), indent=2, sort_keys=True, ) ) return 0 if args.command == "list-jobs": print(json.dumps(list_jobs(root), indent=2, sort_keys=True)) return 0 if args.command == "create-job": print( json.dumps( create_job(root, args.job_name, args.job_kind, args.description), indent=2, sort_keys=True, ) ) return 0 if args.command == "add-job-output": print( json.dumps( add_job_output( root, args.job_name, args.output_name, args.value_type, bound_custom_field=args.bound_custom_field, description=args.description, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "list-job-versions": print(json.dumps(list_job_versions(root, args.job_name), indent=2, sort_keys=True)) return 0 if args.command == "create-job-version": print( json.dumps( create_job_version( root, args.job_name, instruction=args.instruction, capability=args.capability, provider=args.provider, model=args.model, input_basis=args.input_basis, response_schema_json=args.response_schema_json, parameters_json=args.parameters_json, segment_profile=args.segment_profile, aggregation_strategy=args.aggregation_strategy, display_name=args.display_name, ), indent=2, sort_keys=True, ) ) return 0 if args.command == "list-fields": payload = list_fields(root) if args.format == "table": sys.stdout.write(render_list_fields_table(payload) + "\n") sys.stdout.flush() return 0 return emit_cli_payload("list-fields", payload) if args.command == "add-field": return emit_cli_payload("add-field", add_field(root, args.field_name, args.field_type, args.instruction)) if args.command == "rename-field": return emit_cli_payload("rename-field", rename_field(root, args.old_name, args.new_name)) if args.command == "delete-field": return emit_cli_payload("delete-field", delete_field(root, args.field_name, confirm=args.confirm)) if args.command == "describe-field": return emit_cli_payload( "describe-field", describe_field(root, args.field_name, text=args.text, clear=args.clear), ) if args.command == "change-field-type": return emit_cli_payload( "change-field-type", change_field_type(root, args.field_name, args.target_field_type), ) if args.command == "promote-field-type": return emit_cli_payload("promote-field-type", promote_field_type(root, args.field_name, args.target_field_type)) if args.command == "fill-field": return emit_cli_payload( "fill-field", fill_field( root, field_name=args.field, value=args.value, clear=args.clear, document_ids=args.document_ids, query=args.query, raw_bates=args.bates, raw_filters=args.filters, dataset_names=args.dataset_names, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, dry_run=args.dry_run, confirm=args.confirm, ), ) if args.command == "set-field": return emit_cli_payload("set-field", set_field(root, args.document_ids, args.field, args.value)) if args.command == "refresh-text-derived-metadata": return emit_cli_payload( "refresh-text-derived-metadata", refresh_text_derived_metadata( root, document_ids=args.document_ids, query=args.query, raw_bates=args.bates, raw_filters=args.filters, dataset_names=args.dataset_names, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, field_names=args.field_names, ), ) if args.command == "reconcile-duplicates": return emit_cli_payload( "reconcile-duplicates", reconcile_duplicates( root, basis=args.basis, apply_changes=bool(args.apply), document_ids=args.document_ids, query=args.query, raw_bates=args.bates, raw_filters=args.filters, dataset_names=args.dataset_names, from_run_id=args.from_run_id, select_from_scope=args.select_from_scope, ), ) if args.command == "merge-into-conversation": return emit_cli_payload( "merge-into-conversation", merge_into_conversation(root, args.doc_id, args.target_doc_id), ) if args.command == "split-from-conversation": return emit_cli_payload( "split-from-conversation", split_from_conversation(root, args.doc_id), ) if args.command == "clear-conversation-assignment": return emit_cli_payload( "clear-conversation-assignment", clear_conversation_assignment(root, args.doc_id), ) if args.command in {"refresh-conversation-previews", "refresh-previews"}: return emit_cli_payload( args.command, refresh_generated_previews( root, scope=args.scope, conversation_ids=args.conversation_ids, document_ids=args.document_ids, dataset_id=args.dataset_id, dataset_name=args.dataset_name, missing_only=args.missing_only, from_source=args.from_source, ), ) if args.command == "rebuild-conversations": return emit_cli_payload( "rebuild-conversations", rebuild_conversations( root, conversation_ids=args.conversation_ids, document_ids=args.document_ids, dataset_id=args.dataset_id, dataset_name=args.dataset_name, batch_size=args.batch_size, ), ) parser.error(f"Unknown command: {args.command}") return 2 except RetrieverStructuredError as exc: if CLI_OUTPUT_MODE == "human": message = normalize_inline_whitespace(str(exc.payload.get("error") or exc.payload.get("message") or "Retriever error")) print(message, file=sys.stderr) else: print(json.dumps({"tool_version": TOOL_VERSION, **exc.payload}, indent=2, sort_keys=True), file=sys.stderr) return 2 except RetrieverError as exc: if CLI_OUTPUT_MODE == "human": print(str(exc), file=sys.stderr) else: print(json.dumps({"error": str(exc), "tool_version": TOOL_VERSION}), file=sys.stderr) return 2 except sqlite3.Error as exc: message = f"SQLite error: {exc}" if CLI_OUTPUT_MODE == "human": print(message, file=sys.stderr) else: print(json.dumps({"error": message, "tool_version": TOOL_VERSION}), file=sys.stderr) return 2 except Exception as exc: message = f"{type(exc).__name__}: {exc}" if CLI_OUTPUT_MODE == "human": print(message, file=sys.stderr) else: print(json.dumps({"error": message, "tool_version": TOOL_VERSION}), file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())