Merge pull request #171 from ggozad/feat/rebuild-options

New options for `rebuild` command to control what gets rebuilt

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

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
+Rebuild the database by re-indexing documents. Useful when switching embeddings provider/model or changing chunking settings:
-when want to switch embeddings provider or model:
 
 ```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:

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:

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(

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")

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]:
+    async def rebuild_database(
-        """Rebuild the database by deleting all chunks and re-indexing all documents.
+        self, mode: RebuildMode = RebuildMode.FULL
+    ) -> AsyncGenerator[str, None]:
+        """Rebuild the database with the specified mode.
 
-        For documents with URIs:
+        Args:
-        - Re-adds from source if source exists
+            mode: The rebuild mode to use:
-        - Re-embeds from existing content if source is missing
+                - FULL: Re-convert from source files, re-chunk, re-embed (default)
-
+                - RECHUNK: Re-chunk from existing content, re-embed (no source access)
-        For documents without URIs:
+                - EMBED_ONLY: Keep existing chunks, only regenerate embeddings
-        - Re-creates chunks from existing content
 
         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:
+        if mode == RebuildMode.EMBED_ONLY:
-            assert doc.id is not None, "Document ID should not be None"
+            async for doc_id in self._rebuild_embed_only(documents):
-            if doc.uri:
+                yield doc_id
-                # Document has a URI - check if source is accessible
+        elif mode == RebuildMode.RECHUNK:
-                source_accessible = False
+            await self.chunk_repository.delete_all()
-                parsed_url = urlparse(doc.uri)
+            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:
+        # Final maintenance
-                    if parsed_url.scheme == "file":
+        try:
-                        # Check if file exists
+            await self.store.vacuum()
-                        source_path = Path(parsed_url.path)
+        except Exception:
-                        source_accessible = source_path.exists()
+            pass
-                    elif parsed_url.scheme in ("http", "https"):
+
-                        # For URLs, we'll try to create and catch errors
+    async def _rebuild_embed_only(
-                        source_accessible = True
+        self, documents: list[Document]
-                    else:
+    ) -> AsyncGenerator[str, None]:
-                        source_accessible = False
+        """Re-embed all chunks without changing chunk boundaries."""
-                except Exception:
+        for doc in documents:
-                    source_accessible = False
+            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, (
+                        assert new_doc.id is not None
-                            "New document ID should not be 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."""

tests/test_rebuild.py

@@ -1,86 +1,92 @@
 import pytest
 from datasets import Dataset
 
-from haiku.rag.client import HaikuRAG
+from haiku.rag.client import HaikuRAG, RebuildMode
-from haiku.rag.store.models.document import Document
 
 
 @pytest.mark.asyncio
-async def test_rebuild_database(qa_corpus: Dataset, temp_db_path):
+async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
-    """Test rebuild functionality with existing documents."""
+    """Test full rebuild: converts, chunks, and embeds all documents."""
     async with HaikuRAG(temp_db_path) as client:
-        created_docs: list[Document] = []
+        doc = await client.create_document(content=qa_corpus["document_extracted"][0])
-        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
         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_ids = [doc_id async for doc_id in client.rebuild_database()]
-        processed_doc_ids = []
-        async for doc_id in client.rebuild_database():
-            processed_doc_ids.append(doc_id)
 
-        # Document should still be processed (not skipped)
+        assert doc.id in processed_ids
-        assert doc.id in processed_doc_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)