From cb02a2c3d8a22c1384f5f53f6e7d3eb3ff24cb85 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 24 Apr 2026 12:17:21 +0300 Subject: [PATCH] Extract rebuild machinery into client/rebuild.py --- haiku_rag_slim/haiku/rag/client/__init__.py | 303 +----------------- haiku_rag_slim/haiku/rag/client/rebuild.py | 323 ++++++++++++++++++++ 2 files changed, 326 insertions(+), 300 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/client/rebuild.py diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 11853794..d33cdb8c 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -915,307 +915,10 @@ class HaikuRAG: async def rebuild_database( self, mode: RebuildMode = RebuildMode.FULL ) -> AsyncGenerator[str, None]: - """Rebuild the database with the specified mode. + from haiku.rag.client.rebuild import rebuild_database - Args: - mode: The rebuild mode to use: - - FULL: Re-convert from source files, re-chunk, re-embed (default) - - RECHUNK: Re-chunk from existing content, re-embed (no source access) - - EMBED_ONLY: Keep existing chunks, only regenerate embeddings - - TITLE_ONLY: Only generate titles for untitled documents - - Yields: - The ID of the document currently being processed. - """ - # Wait for any background vacuum before destructive table operations - await self._await_vacuum_tasks() - - # Update settings to current config - settings_repo = SettingsRepository(self.store) - await settings_repo.save_current_settings() - - documents = await self.list_documents(include_content=True) - - if mode == RebuildMode.TITLE_ONLY: - async for doc_id in self._rebuild_title_only(documents): - yield doc_id - elif mode == RebuildMode.EMBED_ONLY: - async for doc_id in self._rebuild_embed_only(documents): - yield doc_id - elif mode == RebuildMode.RECHUNK: - await self.chunk_repository.delete_all() - await self.store.recreate_embeddings_table() - async for doc_id in self._rebuild_rechunk(documents): - yield doc_id - else: # FULL - await self.chunk_repository.delete_all() - await self.store.recreate_embeddings_table() - async for doc_id in self._rebuild_full(documents): - yield doc_id - - # Final maintenance if auto_vacuum enabled - if self._config.storage.auto_vacuum: - try: - await self.store.vacuum() - except Exception: - pass - - async def _rebuild_title_only( - self, documents: list[Document] - ) -> AsyncGenerator[str, None]: - """Generate titles for documents that don't have one.""" - for doc in documents: - if doc.title is not None: - continue - assert doc.id is not None - try: - title = await self.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 self.document_repository.update(doc) - yield doc.id - - async def _rebuild_embed_only( - self, documents: list[Document] - ) -> AsyncGenerator[str, None]: - """Re-embed all chunks without changing chunk boundaries.""" - from haiku.rag.embeddings import contextualize - - # Collect all chunks with new embeddings - all_chunk_data: list[tuple[str, dict]] = [] - - for doc in documents: - assert doc.id is not None - chunks = await self.chunk_repository.get_by_document_id(doc.id) - if not chunks: - continue - - texts = contextualize(chunks) - embeddings = await self.chunk_repository.embedder.embed_documents(texts) - - for chunk, content_fts, embedding in zip(chunks, texts, embeddings): - all_chunk_data.append( - ( - doc.id, - { - "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, - }, - ) - ) - - # Recreate chunks table (handles dimension changes) - await self.store.recreate_embeddings_table() - - # Insert all chunks - if all_chunk_data: - records = [self.store.ChunkRecord(**data) for _, data in all_chunk_data] - await self.store.chunks_table.add(records) - - # Yield all processed doc IDs - yielded_docs: set[str] = set() - for doc_id, _ in all_chunk_data: - if doc_id not in yielded_docs: - yielded_docs.add(doc_id) - yield doc_id - - # 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( - self, documents: list[Document], chunks: list[Chunk] - ) -> None: - """Batch write documents and chunks during rebuild. - - This performs two writes: one for all document updates, 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 ( - self.store.documents_table.merge_insert("id") - .when_matched_update_all() - .execute(doc_records) - ) - - # Batch create all chunks (single LanceDB version) - if chunks: - await self.chunk_repository.create(chunks) - - # Repopulate document items from stored docling data - for doc in documents: - assert doc.id is not None - docling_doc = doc.get_docling_document() - if docling_doc is not None: - await self.document_item_repository.delete_by_document_id(doc.id) - items = extract_items(doc.id, docling_doc) - await self.document_item_repository.create_items(doc.id, items) - - async def _rebuild_rechunk( - self, documents: list[Document] - ) -> AsyncGenerator[str, None]: - """Re-chunk and re-embed from existing document content.""" - from haiku.rag.embeddings import embed_chunks - - batch_size = 50 - pending_chunks: list[Chunk] = [] - pending_docs: list[Document] = [] - pending_doc_ids: list[str] = [] - - converter = get_converter(self._config) - - for doc in documents: - assert doc.id is not None - - # Convert stored markdown to DoclingDocument - docling_document = await converter.convert_text(doc.content, format="md") - - # Chunk and embed - chunks = await self.chunk(docling_document) - embedded_chunks = await embed_chunks(chunks, self._config) - - # Update document fields - 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) - pending_doc_ids.append(doc.id) - - # Flush batch when size reached - if len(pending_docs) >= batch_size: - await self._flush_rebuild_batch(pending_docs, pending_chunks) - for doc_id in pending_doc_ids: - yield doc_id - pending_chunks = [] - pending_docs = [] - pending_doc_ids = [] - - # Flush remaining - if pending_docs: - await self._flush_rebuild_batch(pending_docs, pending_chunks) - for doc_id in pending_doc_ids: - yield doc_id - - async def _rebuild_full( - self, documents: list[Document] - ) -> AsyncGenerator[str, None]: - """Full rebuild: re-convert from source, re-chunk, re-embed.""" - from haiku.rag.embeddings import embed_chunks - - batch_size = 50 - pending_chunks: list[Chunk] = [] - pending_docs: list[Document] = [] - pending_doc_ids: list[str] = [] - converter = get_converter(self._config) - - for doc in documents: - assert doc.id is not None - - # Try to rebuild from source if available - if doc.uri and self._check_source_accessible(doc.uri): - try: - # Flush pending batch before source rebuild (creates new doc) - if pending_docs: - await self._flush_rebuild_batch(pending_docs, pending_chunks) - for doc_id in pending_doc_ids: - yield doc_id - pending_chunks = [] - pending_docs = [] - pending_doc_ids = [] - - await self.delete_document(doc.id) - new_doc = await self.create_document_from_source( - source=doc.uri, metadata=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", - doc.uri, - e, - ) - continue - - # Fallback: rebuild from stored content - 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 self.chunk(docling_document) - embedded_chunks = await embed_chunks(chunks, self._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) - pending_doc_ids.append(doc.id) - - # Flush batch when size reached - if len(pending_docs) >= batch_size: - await self._flush_rebuild_batch(pending_docs, pending_chunks) - for doc_id in pending_doc_ids: - yield doc_id - pending_chunks = [] - pending_docs = [] - pending_doc_ids = [] - - # Flush remaining - if pending_docs: - await self._flush_rebuild_batch(pending_docs, pending_chunks) - for doc_id in pending_doc_ids: - yield doc_id + async for doc_id in rebuild_database(self, mode): + yield doc_id def _check_source_accessible(self, uri: str) -> bool: """Check if a document's source URI is accessible.""" diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py new file mode 100644 index 00000000..f9b9e8f5 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -0,0 +1,323 @@ +import json +import logging +from collections.abc import AsyncGenerator +from datetime import datetime +from typing import TYPE_CHECKING + +from haiku.rag.converters import get_converter +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__) + + +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 + + # Wait for any background vacuum before destructive table operations + await client._await_vacuum_tasks() + + # Update settings to current config + settings_repo = SettingsRepository(client.store) + await settings_repo.save_current_settings() + + documents = await client.list_documents(include_content=True) + + 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): + 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 + 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 + if client._config.storage.auto_vacuum: + try: + await client.store.vacuum() + except Exception: + pass + + +async def _rebuild_title_only( + client: "HaikuRAG", documents: list[Document] +) -> AsyncGenerator[str, None]: + """Generate titles for documents that don't have one.""" + for doc in documents: + if doc.title is not None: + continue + assert doc.id is not None + 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) + yield doc.id + + +async def _rebuild_embed_only( + client: "HaikuRAG", documents: list[Document] +) -> AsyncGenerator[str, None]: + """Re-embed all chunks without changing chunk boundaries.""" + from haiku.rag.embeddings import contextualize + + # Collect all chunks with new embeddings + all_chunk_data: list[tuple[str, dict]] = [] + + for doc in documents: + assert doc.id is not None + chunks = await client.chunk_repository.get_by_document_id(doc.id) + if not chunks: + continue + + texts = contextualize(chunks) + embeddings = await client.chunk_repository.embedder.embed_documents(texts) + + for chunk, content_fts, embedding in zip(chunks, texts, embeddings): + all_chunk_data.append( + ( + doc.id, + { + "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, + }, + ) + ) + + # Recreate chunks table (handles dimension changes) + await client.store.recreate_embeddings_table() + + # Insert all chunks + if all_chunk_data: + records = [client.store.ChunkRecord(**data) for _, data in all_chunk_data] + await client.store.chunks_table.add(records) + + # Yield all processed doc IDs + yielded_docs: set[str] = set() + for doc_id, _ in all_chunk_data: + if doc_id not in yielded_docs: + yielded_docs.add(doc_id) + yield doc_id + + # 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 + for doc in documents: + assert doc.id is not None + docling_doc = doc.get_docling_document() + if docling_doc is not None: + await client.document_item_repository.delete_by_document_id(doc.id) + items = extract_items(doc.id, docling_doc) + 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 from existing document content.""" + from haiku.rag.embeddings import embed_chunks + + batch_size = 50 + pending_chunks: list[Chunk] = [] + pending_docs: list[Document] = [] + pending_doc_ids: list[str] = [] + + converter = get_converter(client._config) + + for doc in documents: + assert doc.id is not None + + # Convert stored markdown to DoclingDocument + docling_document = await converter.convert_text(doc.content, format="md") + + # Chunk and embed + chunks = await client.chunk(docling_document) + embedded_chunks = await embed_chunks(chunks, client._config) + + # Update document fields + 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) + pending_doc_ids.append(doc.id) + + # Flush batch when size reached + if len(pending_docs) >= batch_size: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id + pending_chunks = [] + pending_docs = [] + pending_doc_ids = [] + + # Flush remaining + if pending_docs: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id + + +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 + + batch_size = 50 + pending_chunks: list[Chunk] = [] + pending_docs: list[Document] = [] + pending_doc_ids: list[str] = [] + converter = get_converter(client._config) + + for doc in documents: + assert doc.id is not None + + # Try to rebuild from source if available + if doc.uri and client._check_source_accessible(doc.uri): + try: + # Flush pending batch before source rebuild (creates new doc) + if pending_docs: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id + pending_chunks = [] + pending_docs = [] + pending_doc_ids = [] + + await client.delete_document(doc.id) + new_doc = await client.create_document_from_source( + source=doc.uri, metadata=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", + doc.uri, + e, + ) + continue + + # Fallback: rebuild from stored content + 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) + pending_doc_ids.append(doc.id) + + # Flush batch when size reached + if len(pending_docs) >= batch_size: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id + pending_chunks = [] + pending_docs = [] + pending_doc_ids = [] + + # Flush remaining + if pending_docs: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id