import json import logging from collections.abc import AsyncGenerator from datetime import datetime from typing import TYPE_CHECKING from lancedb.pydantic import LanceModel from haiku.rag.client.documents import check_source_accessible from haiku.rag.converters import get_converter from haiku.rag.store.engine import ChunkRecordBase from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.repositories.settings import SettingsRepository if TYPE_CHECKING: from haiku.rag.client import HaikuRAG, RebuildMode logger = logging.getLogger(__name__) _REBUILD_BATCH_SIZE = 50 _STAGING_TABLE_NAME = "chunks_rebuild_staging" _STAGING_MARKER_TABLE_NAME = "chunks_rebuild_marker" _STAGING_COPY_BATCH_SIZE = 1000 class _StagingChunkRecord(LanceModel): """Non-vector copy of a chunk row, used by ``_rebuild_embed_only``. The staging table holds the original chunks' identity and content while the live ``chunks`` table is dropped and recreated with a potentially different vector dimension. The vector itself is omitted — re-embedding is the whole point — and ``content_fts`` is regenerated by ``contextualize`` during phase 2. Mirrors ``ChunkRecordBase`` minus ``content_fts`` and ``vector``. Keep in sync: ``test_staging_chunk_record_mirrors_chunk_record_schema`` enforces parity so a new column on ``ChunkRecordBase`` can't silently get dropped on every embed-only rebuild. """ id: str document_id: str content: str metadata: str order: int class _StagingMarkerRecord(LanceModel): """Sentinel marking the staging table as complete. The marker table is created only after ``_populate_staging_table`` writes every chunk into staging. Its presence at the top of a rebuild means phase 2 (the embed loop) was interrupted by an earlier crash, so staging is the authoritative source for the original chunk identities and we must resume from it instead of rerunning phase 1. """ id: str async def rebuild_database( client: "HaikuRAG", mode: "RebuildMode | None" = None ) -> AsyncGenerator[str, None]: """Rebuild the database with the specified mode. Yields the ID of each document as it is processed. """ from haiku.rag.client import RebuildMode if mode is None: mode = RebuildMode.FULL # Resolve any leftover staging/marker tables from a previously # interrupted rebuild. Returns True only when phase 1 was already # complete and the current mode is EMBED_ONLY, in which case we resume # phase 2 from the existing staging table instead of recopying. resume_from_staging = await _resolve_rebuild_recovery(client, mode) # Wait for any already-scheduled background vacuum before the destructive # table operations at the top of RECHUNK / FULL. Rebuild drops and # recreates tables (and creates indices); a concurrent optimize on the # same table fails with "CreateIndex transaction was preempted" from # lance. Note: FULL calls create_document_from_source inside its loop, # which may schedule *new* background vacuums — those run after the # destructive phase and are fine. await client._await_vacuum_tasks() # Update settings to current config settings_repo = SettingsRepository(client.store) await settings_repo.save_current_settings() # Light listing — id/uri/title/metadata only. Each rebuild function # fetches full content (including the multi-MB docling_pages blob) one # document at a time so a 1000-doc database doesn't pull ~15 GB of # blobs into memory before the loop starts. documents = await client.list_documents(include_content=False) if mode == RebuildMode.TITLE_ONLY: async for doc_id in _rebuild_title_only(client, documents): yield doc_id elif mode == RebuildMode.EMBED_ONLY: async for doc_id in _rebuild_embed_only( client, documents, resume_from_staging=resume_from_staging ): yield doc_id elif mode == RebuildMode.RECHUNK: await client.chunk_repository.delete_all() await client.store.recreate_embeddings_table() async for doc_id in _rebuild_rechunk(client, documents): yield doc_id elif mode == RebuildMode.DESCRIPTIONS: await client.chunk_repository.delete_all() await client.store.recreate_embeddings_table() async for doc_id in _rebuild_descriptions(client, documents): yield doc_id else: # FULL await client.chunk_repository.delete_all() await client.store.recreate_embeddings_table() async for doc_id in _rebuild_full(client, documents): yield doc_id # Final maintenance if auto_vacuum enabled. Swallowing only so that a # failed post-rebuild optimize doesn't mask a successful rebuild — but # log it so the failure is visible in the output. if client._config.storage.auto_vacuum: try: await client.store.vacuum() except Exception: logger.warning("Post-rebuild vacuum failed", exc_info=True) async def _hydrate( client: "HaikuRAG", light_docs: list[Document] ) -> AsyncGenerator[Document, None]: """Yield fully-loaded documents one at a time from a light listing. The light listing in ``rebuild_database`` skips the multi-MB ``docling_document``/``docling_pages`` blobs; this helper fetches each full record on demand so peak memory stays at ~one document. Documents that disappeared between listing and processing are silently skipped. """ for light_doc in light_docs: assert light_doc.id is not None doc = await client.get_document_by_id(light_doc.id) if doc is None: continue assert doc.id is not None yield doc async def _rebuild_title_only( client: "HaikuRAG", documents: list[Document] ) -> AsyncGenerator[str, None]: """Generate titles for documents that don't have one.""" untitled = [d for d in documents if d.title is None] async for doc in _hydrate(client, untitled): try: title = await client.generate_title(doc) except Exception: logger.warning( "Failed to generate title for document %s", doc.id, exc_info=True ) continue if title is not None: doc.title = title await client.document_repository.update(doc) assert doc.id is not None yield doc.id async def _resolve_rebuild_recovery(client: "HaikuRAG", mode: "RebuildMode") -> bool: """Resolve any partially-completed rebuild state from a previous crash. Returns ``True`` if ``_rebuild_embed_only`` should resume from the existing staging table (phase 1 was already complete). In all other cases stale recovery tables are dropped and the rebuild starts fresh. State at entry → action -------------------------------- no staging, no marker → return False (normal start) staging only → drop staging (phase 1 was interrupted; ``chunks`` is intact) marker only → drop marker (corrupted state) staging + marker, embed → return True (resume phase 2 from staging) staging + marker, other → drop both (staging is for embed-only; user picked a different mode) """ from haiku.rag.client import RebuildMode db = client.store.db tables = (await db.list_tables()).tables has_staging = _STAGING_TABLE_NAME in tables has_marker = _STAGING_MARKER_TABLE_NAME in tables if not has_staging and not has_marker: return False if has_marker and not has_staging: logger.warning( "Found '%s' without staging table; dropping orphaned marker.", _STAGING_MARKER_TABLE_NAME, ) await db.drop_table(_STAGING_MARKER_TABLE_NAME) return False if not has_marker: logger.warning( "Dropping incomplete '%s' from an interrupted phase 1.", _STAGING_TABLE_NAME, ) await db.drop_table(_STAGING_TABLE_NAME) return False # has_staging and has_marker if mode == RebuildMode.EMBED_ONLY: logger.warning( "Resuming interrupted embed-only rebuild: phase 2 will run from " "existing '%s'.", _STAGING_TABLE_NAME, ) return True logger.warning( "Dropping staging tables from a prior embed-only rebuild — current " "mode (%s) does not consume them.", mode.name, ) await db.drop_table(_STAGING_MARKER_TABLE_NAME) await db.drop_table(_STAGING_TABLE_NAME) return False async def _populate_staging_table(client: "HaikuRAG") -> None: """Stream the non-vector columns of the chunks table into staging. Uses ``to_batches`` for a single streaming read (no offset/limit pagination drift), so peak memory stays bounded regardless of corpus size. The vector column is omitted — the point of embed-only rebuild is to regenerate it. """ db = client.store.db tables = (await db.list_tables()).tables if _STAGING_TABLE_NAME in tables: await db.drop_table(_STAGING_TABLE_NAME) staging = await db.create_table(_STAGING_TABLE_NAME, schema=_StagingChunkRecord) if "chunks" not in tables: return stream = ( await client.store.chunks_table.query() .select(["id", "document_id", "content", "metadata", "order"]) .to_batches(max_batch_length=_STAGING_COPY_BATCH_SIZE) ) async for batch in stream: rows = batch.to_pylist() if not rows: continue records = [ _StagingChunkRecord( id=r["id"], document_id=r["document_id"], content=r["content"], metadata=r["metadata"], order=r["order"], ) for r in rows ] await staging.add(records) async def _mark_phase1_complete(client: "HaikuRAG") -> None: """Create the marker table that designates staging as authoritative. Called after ``_populate_staging_table`` finishes. On crash recovery the marker's presence flips ``_rebuild_embed_only`` into resume mode. """ db = client.store.db if _STAGING_MARKER_TABLE_NAME in (await db.list_tables()).tables: return marker = await db.create_table( _STAGING_MARKER_TABLE_NAME, schema=_StagingMarkerRecord ) await marker.add([_StagingMarkerRecord(id="phase1_complete")]) async def _drop_staging_tables(client: "HaikuRAG") -> None: """Drop the marker first, then the staging table. Ordering matters: if a crash interrupts cleanup between the two drops, the next rebuild sees ``staging`` without ``marker`` and treats it as a partial phase 1 → drops staging harmlessly. The reverse order would leak a marker pointing at nothing. """ db = client.store.db tables = (await db.list_tables()).tables if _STAGING_MARKER_TABLE_NAME in tables: await db.drop_table(_STAGING_MARKER_TABLE_NAME) if _STAGING_TABLE_NAME in tables: await db.drop_table(_STAGING_TABLE_NAME) async def _read_chunks_from_staging(staging_table, document_id: str) -> list[Chunk]: """Read chunks for one document from the staging table. Only non-vector columns are selected: the staging table may have a different vector dimension than the new chunks table (during a dim migration), and we re-embed anyway. """ rows = ( await staging_table.query() .where(f"document_id = '{document_id}'") .select(["id", "document_id", "content", "metadata", "order"]) .to_arrow() ).to_pylist() chunks: list[Chunk] = [] for row in rows: chunks.append( Chunk( id=row["id"], document_id=row["document_id"], content=row["content"], metadata=json.loads(row["metadata"]), order=row["order"], ) ) chunks.sort(key=lambda c: c.order) return chunks async def _rebuild_embed_only( client: "HaikuRAG", documents: list[Document], *, resume_from_staging: bool = False, ) -> AsyncGenerator[str, None]: """Re-embed all chunks without changing chunk boundaries. Two-phase pattern that keeps peak memory bounded regardless of corpus size and is idempotent across crashes: 1. Stream-copy the chunks table's non-vector columns into a staging table, then write a marker row that designates staging as complete. LanceDB OSS does not support ``rename_table``, so the staging copy is the only safe way to preserve chunk identity while the live ``chunks`` table is dropped and recreated. 2. Drop-and-recreate ``chunks`` with the current schema, then stream from staging one document at a time, re-embed in batches of ``embeddings.batch_size``, and flush to the new chunks table every ``_REBUILD_BATCH_SIZE`` documents. Cleanup runs only on success: a crash anywhere in phase 2 leaves both staging and marker in place so the next rebuild can re-enter phase 2 via ``resume_from_staging=True``. The order of the success cleanup — drop marker before staging — keeps an interruption between the two drops recoverable: the next rebuild sees staging without marker and treats it as a partial phase 1, which is harmless because phase 2 has already finished writing the new chunks table. """ from haiku.rag.embeddings import contextualize db = client.store.db batch_size = client._config.embeddings.batch_size if not resume_from_staging: # Phase 1: copy chunks into staging, then mark it complete. After the # marker exists, a crash will resume phase 2 from staging. await _populate_staging_table(client) await _mark_phase1_complete(client) # Recreate the chunks table fresh (idempotent; handles vector-dim # changes and discards any partial new chunks from a prior crashed # phase 2). await client.store.recreate_embeddings_table() staging_table = await db.open_table(_STAGING_TABLE_NAME) pending_records: list[ChunkRecordBase] = [] yielded_docs: set[str] = set() for doc in documents: assert doc.id is not None chunks = await _read_chunks_from_staging(staging_table, doc.id) if not chunks: continue texts = contextualize(chunks) embeddings: list[list[float]] = [] for i in range(0, len(texts), batch_size): batch_embeddings = await client.chunk_repository.embedder.embed_documents( texts[i : i + batch_size] ) embeddings.extend(batch_embeddings) for chunk, content_fts, embedding in zip(chunks, texts, embeddings): assert chunk.id is not None assert chunk.document_id is not None pending_records.append( client.store.ChunkRecord( id=chunk.id, document_id=chunk.document_id, content=chunk.content, content_fts=content_fts, metadata=json.dumps(chunk.metadata), order=chunk.order, vector=embedding, ) ) yielded_docs.add(doc.id) # Yield per-doc for progress reporting; the actual write batches up # to _REBUILD_BATCH_SIZE docs. If the process is interrupted between # yield and the next flush, the next rebuild resumes phase 2 from # the staging table and redoes the batch (see _rebuild_rechunk for # the original comment on the yield/flush gap). yield doc.id if len(yielded_docs) % _REBUILD_BATCH_SIZE == 0 and pending_records: await client.store.chunks_table.add(pending_records) pending_records = [] if pending_records: await client.store.chunks_table.add(pending_records) # Phase 2 finished. Drop the recovery state — marker first so a crash # between the two drops leaves only staging behind, which the next # rebuild discards harmlessly. await _drop_staging_tables(client) # Yield docs with no chunks for doc in documents: if doc.id and doc.id not in yielded_docs: yield doc.id async def _flush_rebuild_batch( client: "HaikuRAG", documents: list[Document], chunks: list[Chunk] ) -> None: """Batch write documents and chunks during rebuild. Performs two writes: one for all document updates (via merge_insert), one for all chunks. Also repopulates document items from the stored docling document. Used by RECHUNK and FULL modes after the chunks table has been cleared. """ from haiku.rag.store.engine import DocumentRecord if not documents: return now = datetime.now().isoformat() # Batch update documents using merge_insert (single LanceDB version) doc_records = [] for doc in documents: assert doc.id is not None doc_records.append( DocumentRecord( id=doc.id, content=doc.content, uri=doc.uri, title=doc.title, metadata=json.dumps(doc.metadata), docling_document=doc.docling_document, docling_pages=doc.docling_pages, docling_version=doc.docling_version, created_at=doc.created_at.isoformat() if doc.created_at else now, updated_at=now, ) ) await ( client.store.documents_table.merge_insert("id") .when_matched_update_all() .execute(doc_records) ) # Batch create all chunks (single LanceDB version) if chunks: await client.chunk_repository.create(chunks) # Repopulate document items from stored docling data. The stored docling # blob has had its picture URIs stripped (compress_docling_split), so # re-extracting from it would lose picture_data — snapshot the existing # bytes per document and merge them back. for doc in documents: assert doc.id is not None docling_doc = doc.get_docling_document() if docling_doc is not None: existing_picture_data = ( await client.document_item_repository.get_all_picture_data(doc.id) ) await client.document_item_repository.delete_by_document_id(doc.id) items = extract_items( doc.id, docling_doc, existing_picture_data=existing_picture_data, ) await client.document_item_repository.create_items(doc.id, items) async def _rebuild_rechunk( client: "HaikuRAG", documents: list[Document] ) -> AsyncGenerator[str, None]: """Re-chunk and re-embed each document from its stored docling blob.""" from haiku.rag.embeddings import embed_chunks, get_embedder pending_chunks: list[Chunk] = [] pending_docs: list[Document] = [] embedder = get_embedder(client._config) async for doc in _hydrate(client, documents): assert doc.id is not None docling_document = doc.get_docling_document() if docling_document is None: raise ValueError( f"Document {doc.id} has no stored docling document; rechunk " "requires it. Run a full rebuild (without --rechunk) instead." ) # Stored blob has stripped picture URIs; pass the snapshot so # build_picture_chunks (inside chunk()) can recover the bytes. existing_picture_data = ( await client.document_item_repository.get_all_picture_data(doc.id) if embedder.supports_images else None ) chunks = await client.chunk( docling_document, existing_picture_data=existing_picture_data, document_id=doc.id, ) embedded_chunks = await embed_chunks(chunks, client._config) # Prepare chunks with document_id and order for order, chunk in enumerate(embedded_chunks): chunk.document_id = doc.id chunk.order = order pending_chunks.extend(embedded_chunks) pending_docs.append(doc) # Yield per-doc so progress reporting moves immediately. The actual # write batches up to _REBUILD_BATCH_SIZE for throughput; if the # process is interrupted between yield and flush, up to one # batch's worth of trailing yields aren't persisted, which is # consistent with the rebuild already being non-atomic. yield doc.id # Flush batch when size reached if len(pending_docs) >= _REBUILD_BATCH_SIZE: await _flush_rebuild_batch(client, pending_docs, pending_chunks) pending_chunks = [] pending_docs = [] # Flush remaining if pending_docs: await _flush_rebuild_batch(client, pending_docs, pending_chunks) async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int: """Run the VLM against pictures lacking a description, patch the docling blob in-place. Returns the number of newly described pictures. Pictures that already carry ``meta.description.text`` are skipped, so the operation is safe to re-run after a partial failure. """ from haiku.rag.providers.picture_description import describe_pictures assert doc.id is not None docling_doc = doc.get_docling_document() if docling_doc is None or not docling_doc.pictures: return 0 needs_description: list[str] = [] for pic in docling_doc.pictures: existing = ( pic.meta.description.text if pic.meta and pic.meta.description else None ) if not (existing and existing.strip()): needs_description.append(pic.self_ref) if not needs_description: return 0 bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk( doc.id, needs_description ) if not bytes_by_ref: logger.warning( "Document %s has %d pictures missing descriptions but no stored " "picture bytes — skipping. Run a full rebuild from source to " "recover the bytes.", doc.id, len(needs_description), ) return 0 descriptions = await describe_pictures(bytes_by_ref, config=client._config) if not descriptions: return 0 # Patch the docling document in-place. PictureMeta + DescriptionMetaField # are pydantic models; build them and assign. from docling_core.types.doc.document import ( DescriptionMetaField, PictureMeta, ) for pic in docling_doc.pictures: text = descriptions.get(pic.self_ref) if not text: continue if pic.meta is None: pic.meta = PictureMeta() pic.meta.description = DescriptionMetaField(text=text) # Update only docling_document — set_docling would also overwrite # docling_pages by routing through compress_docling_split, which # extracts pages from the in-memory JSON and finds none (the pages # blob is stored separately and is not loaded by get_docling_document). # That would silently destroy page rasters for every doc with at # least one undescribed picture. from haiku.rag.store.compression import compress_docling_split structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json()) doc.docling_document = structure_bytes doc.docling_version = docling_doc.version return len(descriptions) async def _rebuild_descriptions( client: "HaikuRAG", documents: list[Document] ) -> AsyncGenerator[str, None]: """Run the VLM over already-stored picture bytes, patch descriptions into the docling blob, then re-chunk + re-embed. Skips the docling parse entirely (the blob is already there); only the VLM cost remains. Idempotent: pictures whose ``meta.description.text`` is already populated are not re-described. """ from haiku.rag.embeddings import embed_chunks, get_embedder if client._config.processing.pictures != "description": raise ValueError( "rebuild --descriptions requires processing.pictures = 'description' " "in your config." ) pending_chunks: list[Chunk] = [] pending_docs: list[Document] = [] embedder = get_embedder(client._config) described_total = 0 async for doc in _hydrate(client, documents): assert doc.id is not None docling_document = doc.get_docling_document() if docling_document is None: raise ValueError( f"Document {doc.id} has no stored docling document; " "rebuild --descriptions requires it. Run a full rebuild instead." ) n = await _patch_picture_descriptions(client, doc) described_total += n # Use the (possibly patched) docling document for chunking. docling_document = doc.get_docling_document() assert docling_document is not None existing_picture_data = ( await client.document_item_repository.get_all_picture_data(doc.id) if embedder.supports_images else None ) chunks = await client.chunk( docling_document, existing_picture_data=existing_picture_data, document_id=doc.id, ) embedded_chunks = await embed_chunks(chunks, client._config) for order, chunk in enumerate(embedded_chunks): chunk.document_id = doc.id chunk.order = order pending_chunks.extend(embedded_chunks) pending_docs.append(doc) yield doc.id if len(pending_docs) >= _REBUILD_BATCH_SIZE: await _flush_rebuild_batch(client, pending_docs, pending_chunks) pending_chunks = [] pending_docs = [] if pending_docs: await _flush_rebuild_batch(client, pending_docs, pending_chunks) logger.info( "rebuild --descriptions: %d new picture descriptions added across %d documents", described_total, len(documents), ) async def _rebuild_full( client: "HaikuRAG", documents: list[Document] ) -> AsyncGenerator[str, None]: """Full rebuild: re-convert from source, re-chunk, re-embed.""" from haiku.rag.embeddings import embed_chunks pending_chunks: list[Chunk] = [] pending_docs: list[Document] = [] converter = get_converter(client._config) for light_doc in documents: assert light_doc.id is not None # Try to rebuild from source if available — uses the light listing # directly, no need to load the stored content/blobs first. if light_doc.uri and check_source_accessible(light_doc.uri): try: # Flush pending batch before source rebuild (creates new doc) if pending_docs: await _flush_rebuild_batch(client, pending_docs, pending_chunks) pending_chunks = [] pending_docs = [] await client.delete_document(light_doc.id) new_doc = await client.create_document_from_source( source=light_doc.uri, metadata=light_doc.metadata or {} ) assert isinstance(new_doc, Document) assert new_doc.id is not None yield new_doc.id continue except Exception as e: logger.error( "Error recreating document from source %s: %s", light_doc.uri, e, ) continue # Fallback: rebuild from stored content. Now we need the full # record (content + docling_pages for the round-trip write). doc = await client.get_document_by_id(light_doc.id) if doc is None: continue assert doc.id is not None if doc.uri: logger.warning("Source missing for %s, re-embedding from content", doc.uri) docling_document = await converter.convert_text(doc.content, format="md") chunks = await client.chunk(docling_document) embedded_chunks = await embed_chunks(chunks, client._config) doc.set_docling(docling_document) # Prepare chunks with document_id and order for order, chunk in enumerate(embedded_chunks): chunk.document_id = doc.id chunk.order = order pending_chunks.extend(embedded_chunks) pending_docs.append(doc) yield doc.id # Flush batch when size reached if len(pending_docs) >= _REBUILD_BATCH_SIZE: await _flush_rebuild_batch(client, pending_docs, pending_chunks) pending_chunks = [] pending_docs = [] # Flush remaining if pending_docs: await _flush_rebuild_batch(client, pending_docs, pending_chunks)