--embed-only, --rechunk options in rebuild command

This commit is contained in:
Yiorgis Gozadinos 2025-11-28 10:11:52 +02:00
parent 9302de0aff
commit 88a4edb537
No known key found for this signature in database
7 changed files with 247 additions and 119 deletions

View file

@ -1,6 +1,14 @@
# Changelog # Changelog
## [Unreleased] ## [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 ## [0.19.3] - 2025-11-27
### Changed ### Changed

View file

@ -259,13 +259,27 @@ haiku-rag vacuum
### Rebuild Database ### 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 ```bash
# Full rebuild (default) - re-converts from source files, re-chunks, re-embeds
haiku-rag rebuild 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 Models
Download required runtime models: Download required runtime models:

View file

@ -157,10 +157,27 @@ await client.delete_document(doc.id)
### Rebuilding the Database ### Rebuilding the Database
```python ```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(): async for doc_id in client.rebuild_database():
print(f"Processed document {doc_id}") 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 ## Maintenance
Run maintenance to optimize storage and prune old table versions: Run maintenance to optimize storage and prune old table versions:

View file

@ -8,7 +8,7 @@ from rich.console import Console
from rich.markdown import Markdown from rich.markdown import Markdown
from rich.progress import Progress 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.config import AppConfig, Config
from haiku.rag.graph.agui import AGUIConsoleRenderer, stream_graph from haiku.rag.graph.agui import AGUIConsoleRenderer, stream_graph
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
@ -406,7 +406,7 @@ class HaikuRAGApp:
except Exception as e: except Exception as e:
self.console.print(f"[red]Error during research: {e}[/red]") 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( async with HaikuRAG(
db_path=self.db_path, config=self.config, skip_validation=True db_path=self.db_path, config=self.config, skip_validation=True
) as client: ) as client:
@ -420,12 +420,18 @@ class HaikuRAGApp:
) )
return return
mode_desc = {
RebuildMode.FULL: "full rebuild",
RebuildMode.RECHUNK: "rechunk",
RebuildMode.EMBED_ONLY: "embed only",
}[mode]
self.console.print( 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: with Progress() as progress:
task = progress.add_task("Rebuilding...", total=total_docs) 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) progress.update(task, advance=1)
self.console.print( self.console.print(

View file

@ -358,9 +358,32 @@ def rebuild(
"--db", "--db",
help="Path to the LanceDB database file", 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) 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") @cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")

View file

@ -3,6 +3,7 @@ import logging
import mimetypes import mimetypes
import tempfile import tempfile
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from enum import Enum
from pathlib import Path from pathlib import Path
from urllib.parse import urlparse from urllib.parse import urlparse
@ -21,6 +22,14 @@ from haiku.rag.store.repositories.settings import SettingsRepository
logger = logging.getLogger(__name__) 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: class HaikuRAG:
"""High-level haiku-rag client.""" """High-level haiku-rag client."""
@ -693,62 +702,103 @@ class HaikuRAG:
) )
return await qa_agent.answer(question) 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: 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 # Update settings to current config
settings_repo = SettingsRepository(self.store) settings_repo = SettingsRepository(self.store)
settings_repo.save_current_settings() settings_repo.save_current_settings()
documents = await self.list_documents() 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: if source_accessible:
# Source exists - delete and recreate from source
try: try:
await self.delete_document(doc.id) await self.delete_document(doc.id)
new_doc = await self.create_document_from_source( new_doc = await self.create_document_from_source(
source=doc.uri, metadata=doc.metadata or {} source=doc.uri, metadata=doc.metadata or {}
) )
# URIs always point to single files/URLs, never directories
assert isinstance(new_doc, Document) 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 yield new_doc.id
except Exception as e: except Exception as e:
logger.error( logger.error(
@ -758,7 +808,6 @@ class HaikuRAG:
) )
continue continue
else: else:
# Source missing - re-embed from existing content
logger.warning( logger.warning(
"Source missing for %s, re-embedding from content", doc.uri "Source missing for %s, re-embedding from content", doc.uri
) )
@ -768,18 +817,23 @@ class HaikuRAG:
) )
yield doc.id yield doc.id
else: else:
# Document without URI - re-create chunks from existing content
docling_document = await converter.convert_text(doc.content) docling_document = await converter.convert_text(doc.content)
await self.chunk_repository.create_chunks_for_document( await self.chunk_repository.create_chunks_for_document(
doc.id, docling_document doc.id, docling_document
) )
yield doc.id 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: 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: except Exception:
pass return False
async def vacuum(self) -> None: async def vacuum(self) -> None:
"""Optimize and clean up old versions across all tables.""" """Optimize and clean up old versions across all tables."""

View file

@ -1,86 +1,92 @@
import pytest import pytest
from datasets import Dataset 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 @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: 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 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) chunks_before = await client.chunk_repository.get_by_document_id(doc.id)
assert len(chunks_before) > 0 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) chunks_after = await client.chunk_repository.get_by_document_id(doc.id)
assert len(chunks_after) > 0 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)