danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -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

18
docs/cli.md
View file

@ -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:

17
docs/python.md
View file

@ -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:

14
haiku_rag_slim/haiku/rag/app.py
View file

@ -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(

25
haiku_rag_slim/haiku/rag/cli.py
View file

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

138
haiku_rag_slim/haiku/rag/client.py
View file

@ -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."""

146
tests/test_rebuild.py
View file

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