diff --git a/CHANGELOG.md b/CHANGELOG.md index d71dedf2..c718ace4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- **Rebuild Modes**: New options for `rebuild` command to control what gets rebuilt + - `--embed-only`: Only regenerate embeddings, keeping existing chunks (fastest option when changing embedding model) + - `--rechunk`: Re-chunk from existing document content without accessing source files + - Default (no flag): Full rebuild with source file re-conversion + - Python API: `rebuild_database(mode=RebuildMode.EMBED_ONLY | RECHUNK | FULL)` + ## [0.19.3] - 2025-11-27 ### Changed diff --git a/docs/cli.md b/docs/cli.md index 690d57b1..4b8b7d68 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -259,13 +259,27 @@ haiku-rag vacuum ### Rebuild Database -Rebuild the database by deleting all chunks & embeddings and re-indexing all documents. This is useful -when want to switch embeddings provider or model: +Rebuild the database by re-indexing documents. Useful when switching embeddings provider/model or changing chunking settings: ```bash +# Full rebuild (default) - re-converts from source files, re-chunks, re-embeds haiku-rag rebuild + +# Re-chunk from stored content (no source file access) +haiku-rag rebuild --rechunk + +# Only regenerate embeddings (fastest, keeps existing chunks) +haiku-rag rebuild --embed-only ``` +**Rebuild modes:** + +| Mode | Flag | Use case | +|------|------|----------| +| Full | (default) | Changed converter, source files updated | +| Rechunk | `--rechunk` | Changed chunking strategy or chunk size | +| Embed only | `--embed-only` | Changed embedding model or vector dimensions | + ### Download Models Download required runtime models: diff --git a/docs/python.md b/docs/python.md index 823a690d..12ac32ab 100644 --- a/docs/python.md +++ b/docs/python.md @@ -157,10 +157,27 @@ await client.delete_document(doc.id) ### Rebuilding the Database ```python +from haiku.rag.client import RebuildMode + +# Full rebuild (default) - re-converts from source files, re-chunks, re-embeds async for doc_id in client.rebuild_database(): print(f"Processed document {doc_id}") + +# Re-chunk from stored content (no source file access) +async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK): + print(f"Processed document {doc_id}") + +# Only regenerate embeddings (fastest, keeps existing chunks) +async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY): + print(f"Processed document {doc_id}") ``` +**Rebuild modes:** + +- `RebuildMode.FULL` - Re-convert from source files, re-chunk, re-embed (default) +- `RebuildMode.RECHUNK` - Re-chunk from existing document content, re-embed +- `RebuildMode.EMBED_ONLY` - Keep existing chunks, only regenerate embeddings + ## Maintenance Run maintenance to optimize storage and prune old table versions: diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 7e75a575..808007c2 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -8,7 +8,7 @@ from rich.console import Console from rich.markdown import Markdown from rich.progress import Progress -from haiku.rag.client import HaikuRAG +from haiku.rag.client import HaikuRAG, RebuildMode from haiku.rag.config import AppConfig, Config from haiku.rag.graph.agui import AGUIConsoleRenderer, stream_graph from haiku.rag.graph.research.dependencies import ResearchContext @@ -406,7 +406,7 @@ class HaikuRAGApp: except Exception as e: self.console.print(f"[red]Error during research: {e}[/red]") - async def rebuild(self): + async def rebuild(self, mode: RebuildMode = RebuildMode.FULL): async with HaikuRAG( db_path=self.db_path, config=self.config, skip_validation=True ) as client: @@ -420,12 +420,18 @@ class HaikuRAGApp: ) return + mode_desc = { + RebuildMode.FULL: "full rebuild", + RebuildMode.RECHUNK: "rechunk", + RebuildMode.EMBED_ONLY: "embed only", + }[mode] + self.console.print( - f"[bold cyan]Rebuilding database with {total_docs} documents...[/bold cyan]" + f"[bold cyan]Rebuilding database ({mode_desc}) with {total_docs} documents...[/bold cyan]" ) with Progress() as progress: task = progress.add_task("Rebuilding...", total=total_docs) - async for _ in client.rebuild_database(): + async for _ in client.rebuild_database(mode=mode): progress.update(task, advance=1) self.console.print( diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index afc5af4e..f8fa9a97 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -358,9 +358,32 @@ def rebuild( "--db", help="Path to the LanceDB database file", ), + embed_only: bool = typer.Option( + False, + "--embed-only", + help="Only regenerate embeddings, keep existing chunks", + ), + rechunk: bool = typer.Option( + False, + "--rechunk", + help="Re-chunk from existing content without accessing source files", + ), ): + from haiku.rag.client import RebuildMode + + if embed_only and rechunk: + typer.echo("Error: --embed-only and --rechunk are mutually exclusive") + raise typer.Exit(1) + + if embed_only: + mode = RebuildMode.EMBED_ONLY + elif rechunk: + mode = RebuildMode.RECHUNK + else: + mode = RebuildMode.FULL + app = create_app(db) - asyncio.run(app.rebuild()) + asyncio.run(app.rebuild(mode=mode)) @cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage") diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 81513c05..c09f2666 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -3,6 +3,7 @@ import logging import mimetypes import tempfile from collections.abc import AsyncGenerator +from enum import Enum from pathlib import Path from urllib.parse import urlparse @@ -21,6 +22,14 @@ from haiku.rag.store.repositories.settings import SettingsRepository logger = logging.getLogger(__name__) +class RebuildMode(Enum): + """Mode for rebuilding the database.""" + + FULL = "full" # Re-convert from source, re-chunk, re-embed + RECHUNK = "rechunk" # Re-chunk from existing content, re-embed + EMBED_ONLY = "embed_only" # Keep chunks, only regenerate embeddings + + class HaikuRAG: """High-level haiku-rag client.""" @@ -693,62 +702,103 @@ class HaikuRAG: ) return await qa_agent.answer(question) - async def rebuild_database(self) -> AsyncGenerator[str, None]: - """Rebuild the database by deleting all chunks and re-indexing all documents. + async def rebuild_database( + self, mode: RebuildMode = RebuildMode.FULL + ) -> AsyncGenerator[str, None]: + """Rebuild the database with the specified mode. - For documents with URIs: - - Re-adds from source if source exists - - Re-embeds from existing content if source is missing - - For documents without URIs: - - Re-creates chunks from existing content + 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 Yields: - int: The ID of the document currently being processed + The ID of the document currently being processed. """ - await self.chunk_repository.delete_all() - self.store.recreate_embeddings_table() - - converter = get_converter(self._config) - # Update settings to current config settings_repo = SettingsRepository(self.store) settings_repo.save_current_settings() documents = await self.list_documents() - for doc in documents: - assert doc.id is not None, "Document ID should not be None" - if doc.uri: - # Document has a URI - check if source is accessible - source_accessible = False - parsed_url = urlparse(doc.uri) + if 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() + 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() + self.store.recreate_embeddings_table() + async for doc_id in self._rebuild_full(documents): + yield doc_id - try: - if parsed_url.scheme == "file": - # Check if file exists - source_path = Path(parsed_url.path) - source_accessible = source_path.exists() - elif parsed_url.scheme in ("http", "https"): - # For URLs, we'll try to create and catch errors - source_accessible = True - else: - source_accessible = False - except Exception: - source_accessible = False + # Final maintenance + try: + await self.store.vacuum() + except Exception: + pass + + async def _rebuild_embed_only( + self, documents: list[Document] + ) -> AsyncGenerator[str, None]: + """Re-embed all chunks without changing chunk boundaries.""" + 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 + + # Batch embed all chunk contents + contents = [chunk.content for chunk in chunks] + embeddings = await self.chunk_repository.embedder.embed(contents) + + # Update each chunk with new embedding + for chunk, embedding in zip(chunks, embeddings): + assert chunk.id is not None + self.store.chunks_table.update( + where=f"id = '{chunk.id}'", + values={"vector": embedding}, + ) + + yield doc.id + + async def _rebuild_rechunk( + self, documents: list[Document] + ) -> AsyncGenerator[str, None]: + """Re-chunk and re-embed from existing document content.""" + converter = get_converter(self._config) + + for doc in documents: + assert doc.id is not None + docling_document = await converter.convert_text(doc.content) + await self.chunk_repository.create_chunks_for_document( + doc.id, docling_document + ) + yield doc.id + + async def _rebuild_full( + self, documents: list[Document] + ) -> AsyncGenerator[str, None]: + """Full rebuild: re-convert from source, re-chunk, re-embed.""" + converter = get_converter(self._config) + + for doc in documents: + assert doc.id is not None + if doc.uri: + source_accessible = self._check_source_accessible(doc.uri) if source_accessible: - # Source exists - delete and recreate from source try: await self.delete_document(doc.id) new_doc = await self.create_document_from_source( source=doc.uri, metadata=doc.metadata or {} ) - # URIs always point to single files/URLs, never directories assert isinstance(new_doc, Document) - assert new_doc.id is not None, ( - "New document ID should not be None" - ) + assert new_doc.id is not None yield new_doc.id except Exception as e: logger.error( @@ -758,7 +808,6 @@ class HaikuRAG: ) continue else: - # Source missing - re-embed from existing content logger.warning( "Source missing for %s, re-embedding from content", doc.uri ) @@ -768,18 +817,23 @@ class HaikuRAG: ) yield doc.id else: - # Document without URI - re-create chunks from existing content docling_document = await converter.convert_text(doc.content) await self.chunk_repository.create_chunks_for_document( doc.id, docling_document ) yield doc.id - # Final maintenance: centralized vacuum to curb disk usage + def _check_source_accessible(self, uri: str) -> bool: + """Check if a document's source URI is accessible.""" + parsed_url = urlparse(uri) try: - await self.store.vacuum() + if parsed_url.scheme == "file": + return Path(parsed_url.path).exists() + elif parsed_url.scheme in ("http", "https"): + return True + return False except Exception: - pass + return False async def vacuum(self) -> None: """Optimize and clean up old versions across all tables.""" diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 0ec81e75..60a04e81 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -1,86 +1,92 @@ import pytest from datasets import Dataset -from haiku.rag.client import HaikuRAG -from haiku.rag.store.models.document import Document +from haiku.rag.client import HaikuRAG, RebuildMode @pytest.mark.asyncio -async def test_rebuild_database(qa_corpus: Dataset, temp_db_path): - """Test rebuild functionality with existing documents.""" +async def test_rebuild_full(qa_corpus: Dataset, temp_db_path): + """Test full rebuild: converts, chunks, and embeds all documents.""" async with HaikuRAG(temp_db_path) as client: - created_docs: list[Document] = [] - for content in qa_corpus["document_extracted"][:3]: - doc = await client.create_document( - content=content, - ) - created_docs.append(doc) - - documents_before = await client.list_documents() - assert len(documents_before) == 3 - - chunks_before = [] - for doc in created_docs: - assert doc.id is not None - doc_chunks = await client.chunk_repository.get_by_document_id(doc.id) - chunks_before.extend(doc_chunks) - - assert len(chunks_before) > 0 - - # Perform rebuild - processed_doc_ids = [] - async for doc_id in client.rebuild_database(): - processed_doc_ids.append(doc_id) - - # Verify all documents were processed - expected_doc_ids = [doc.id for doc in created_docs] - assert set(processed_doc_ids) == set(expected_doc_ids) - - documents_after = await client.list_documents() - assert len(documents_after) == 3 - - # Verify chunks were recreated - chunks_after = [] - for doc in documents_after: - if doc.id is not None: - doc_chunks = await client.chunk_repository.get_by_document_id(doc.id) - chunks_after.extend(doc_chunks) - - assert len(chunks_after) > 0 - - -@pytest.mark.asyncio -async def test_rebuild_with_missing_source(qa_corpus: Dataset, temp_db_path): - """Test rebuild functionality when document source is missing.""" - async with HaikuRAG(temp_db_path) as client: - # Create document with content - content = qa_corpus["document_extracted"][0] - doc = await client.create_document(content=content) - - # Manually set a URI that doesn't exist + doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None - doc_with_uri = await client.document_repository.get_by_id(doc.id) - assert doc_with_uri is not None - doc_with_uri.uri = "file:///nonexistent/path.txt" - await client.document_repository.update(doc_with_uri) - # Verify chunks exist before rebuild chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 + chunk_ids_before = {c.id for c in chunks_before} - # Perform rebuild - processed_doc_ids = [] - async for doc_id in client.rebuild_database(): - processed_doc_ids.append(doc_id) + processed_ids = [doc_id async for doc_id in client.rebuild_database()] - # Document should still be processed (not skipped) - assert doc.id in processed_doc_ids + assert doc.id in processed_ids - # Verify document still exists - doc_after = await client.document_repository.get_by_id(doc.id) - assert doc_after is not None - assert doc_after.content == content - - # Verify chunks were recreated from content chunks_after = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_after) > 0 + chunk_ids_after = {c.id for c in chunks_after} + + # Chunk IDs should change (chunks are recreated) + assert chunk_ids_before.isdisjoint(chunk_ids_after) + + +@pytest.mark.asyncio +async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path): + """Test embed-only rebuild: keeps chunks, only regenerates embeddings.""" + async with HaikuRAG(temp_db_path) as client: + doc = await client.create_document(content=qa_corpus["document_extracted"][0]) + assert doc.id is not None + + chunks_before = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks_before) > 0 + chunk_ids_before = {c.id for c in chunks_before} + chunk_contents_before = {c.id: c.content for c in chunks_before} + + processed_ids = [ + doc_id + async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY) + ] + + assert doc.id in processed_ids + + chunks_after = await client.chunk_repository.get_by_document_id(doc.id) + chunk_ids_after = {c.id for c in chunks_after} + + # Chunk IDs should be preserved (same chunks, just re-embedded) + assert chunk_ids_before == chunk_ids_after + + # Content should be identical + for chunk in chunks_after: + assert chunk.content == chunk_contents_before[chunk.id] + + +@pytest.mark.asyncio +async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path): + """Test rechunk rebuild: re-chunks from content without accessing source files.""" + async with HaikuRAG(temp_db_path) as client: + doc = await client.create_document(content=qa_corpus["document_extracted"][0]) + assert doc.id is not None + + # Set a fake URI to simulate a document that came from a file + doc.uri = "file:///nonexistent/path.txt" + await client.document_repository.update(doc) + + chunks_before = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks_before) > 0 + chunk_ids_before = {c.id for c in chunks_before} + content_before = doc.content + + processed_ids = [ + doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK) + ] + + assert doc.id in processed_ids + + # Document content should be unchanged + doc_after = await client.document_repository.get_by_id(doc.id) + assert doc_after is not None + assert doc_after.content == content_before + + chunks_after = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks_after) > 0 + chunk_ids_after = {c.id for c in chunks_after} + + # Chunk IDs should change (chunks are recreated) + assert chunk_ids_before.isdisjoint(chunk_ids_after)