From 0e3d53791f8ef6553500b7a90c2fb740f7428f67 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 5 Jun 2026 10:42:15 +0300 Subject: [PATCH] Add rebuild --set-embedder to reconcile embedder identity --- CHANGELOG.md | 6 +- docs/cli.md | 6 ++ haiku_rag_slim/haiku/rag/app.py | 8 ++ haiku_rag_slim/haiku/rag/cli.py | 17 ++++- haiku_rag_slim/haiku/rag/client/__init__.py | 2 + haiku_rag_slim/haiku/rag/client/rebuild.py | 27 +++++++ tests/test_rebuild.py | 81 +++++++++++++++++++++ 7 files changed, 143 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 670bef4e..57b116d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,9 +1,13 @@ # Changelog ## [Unreleased] +### Added + +- `haiku-rag rebuild --set-embedder`: adopt the current embedder identity (provider/name) without re-embedding, when the vector dimension is unchanged. + ### Fixed -- Opening a database no longer writes to it: reads no longer rewrite the stored embedding settings or change the stored version, and the version is never downgraded. +- Opening a database no longer writes to it: reads no longer rewrite the stored embedding settings or change the stored version, and the version is never downgraded. Embedding provider/name drift (matching `vector_dim`) warns on read-only opens and raises `ConfigMismatchError` on writable opens; reconcile with `rebuild --set-embedder`. ## [0.54.0] - 2026-06-04 diff --git a/docs/cli.md b/docs/cli.md index 25426a2a..7fe283ad 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -373,6 +373,9 @@ haiku-rag rebuild --title-only # Run the VLM over already-stored picture bytes and patch descriptions # into the docling blob. Skips the docling parse entirely. haiku-rag rebuild --descriptions + +# Adopt the current embedder identity without re-embedding (same vector dimension) +haiku-rag rebuild --set-embedder ``` **Rebuild modes:** @@ -384,6 +387,9 @@ haiku-rag rebuild --descriptions | Embed only | `--embed-only` | Changed embedding model or vector dimensions | | Title only | `--title-only` | Generate titles for documents without one | | Descriptions | `--descriptions` | Add VLM picture descriptions to an existing database | +| Set embedder | `--set-embedder` | Same model, different serving stack (e.g. Ollama to vLLM); vector dimension unchanged | + +**`--set-embedder` mode** updates the stored embedding provider/name to match the current config without re-embedding, valid only when the vector dimension is unchanged. Use it when the same model is served by a different stack so the recorded identity stops drifting from the config. A changed vector dimension is rejected; regenerate embeddings with `--embed-only` or a full rebuild instead. **`--descriptions` mode** runs the configured VLM (`processing.conversion_options.picture_description.model`) over the picture bytes already stored in `document_items.picture_data`, patches each description into the stored docling blob's `pictures[i].meta.description.text`, and re-chunks + re-embeds so chunk text reflects the new descriptions. Requires `processing.pictures: description` in the config. Idempotent: pictures that already carry a description are skipped, so the operation is safe to re-run after a partial failure. The docling parse is skipped entirely. Only the VLM time is paid. diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 185ee07a..485c38d6 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -506,6 +506,14 @@ class HaikuRAGApp: # pragma: no cover read_only=self.read_only, before=self.before, ) as client: + if mode == RebuildMode.SET_EMBEDDER: + async for _ in client.rebuild_database(mode=mode): + pass + self.console.print( + "[bold green]Stored embedder settings updated.[/bold green]" + ) + return + documents = await client.list_documents() total_docs = len(documents) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 26b3dd1c..d646baf2 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -476,14 +476,23 @@ def rebuild( "parse entirely. Requires processing.pictures='description'." ), ), + set_embedder: bool = typer.Option( + False, + "--set-embedder", + help=( + "Adopt the current embedder identity without re-embedding, when the " + "vector dimension is unchanged. Use after swapping the serving stack " + "for the same model (e.g. Ollama to vLLM)." + ), + ), ): from haiku.rag.client import RebuildMode - exclusive = sum([embed_only, rechunk, title_only, descriptions]) + exclusive = sum([embed_only, rechunk, title_only, descriptions, set_embedder]) if exclusive > 1: typer.echo( - "Error: --embed-only, --rechunk, --title-only, and --descriptions " - "are mutually exclusive" + "Error: --embed-only, --rechunk, --title-only, --descriptions, and " + "--set-embedder are mutually exclusive" ) raise typer.Exit(1) @@ -495,6 +504,8 @@ def rebuild( mode = RebuildMode.TITLE_ONLY elif descriptions: # pragma: no cover mode = RebuildMode.DESCRIPTIONS + elif set_embedder: # pragma: no cover + mode = RebuildMode.SET_EMBEDDER else: # pragma: no cover mode = RebuildMode.FULL diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index c4eac92d..61badec5 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -49,6 +49,8 @@ class RebuildMode(Enum): TITLE_ONLY = "title_only" # Only generate titles for untitled documents DESCRIPTIONS = "descriptions" # Run the VLM over already-stored picture # bytes, patch descriptions into the docling blob, then re-chunk + re-embed. + SET_EMBEDDER = "set_embedder" # Adopt the current embedder identity without + # re-embedding, when the vector dimension is unchanged. class HaikuRAG: diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index f44b6b38..25070eb1 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -72,6 +72,10 @@ async def rebuild_database( if mode is None: mode = RebuildMode.FULL + if mode == RebuildMode.SET_EMBEDDER: + await _set_embedder(client) + return + # Resolve any leftover staging/marker tables from a previously # interrupted rebuild. Returns True only when phase 1 was already # complete and the current mode is EMBED_ONLY, in which case we resume @@ -131,6 +135,29 @@ async def rebuild_database( logger.warning("Post-rebuild vacuum failed", exc_info=True) +async def _set_embedder(client: "HaikuRAG") -> None: + """Adopt the current embedder identity without re-embedding. + + Only valid when the vector dimension is unchanged — the stored vectors stay + usable, so just the recorded provider/name are updated. A changed dimension + requires regenerating every embedding via a full rebuild. + """ + from haiku.rag.store.repositories.settings import ConfigMismatchError + + settings_repo = SettingsRepository(client.store) + stored = await settings_repo.get_current_settings() + stored_dim = stored.get("embeddings", {}).get("model", {}).get("vector_dim") + current_dim = client._config.embeddings.model.vector_dim + + if stored_dim is not None and current_dim != stored_dim: + raise ConfigMismatchError( + f"Stored vector dimension {stored_dim} differs from current " + f"{current_dim}; embeddings must be regenerated. Run 'haiku-rag rebuild'." + ) + + await settings_repo.save_current_settings() + + async def _hydrate( client: "HaikuRAG", light_docs: list[Document] ) -> AsyncGenerator[Document, None]: diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index f15aa8bb..1bf9ebd0 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -960,3 +960,84 @@ async def test_rebuild_descriptions_raises_when_blob_is_missing( with pytest.raises(ValueError, match="rebuild --descriptions requires"): async for _ in rag.rebuild_database(mode=RebuildMode.DESCRIPTIONS): pass + + +async def _add_chunk(client: HaikuRAG, vector: list[float]) -> str: + """Insert a chunk row directly, bypassing the embedder.""" + record = client.store.ChunkRecord( + id="chunk-1", + document_id="doc-1", + content="hello", + content_fts="hello", + metadata="{}", + order=0, + vector=vector, + ) + await client.store.chunks_table.add([record]) + return record.id + + +async def _stored_embedding_name(client: HaikuRAG) -> str: + from haiku.rag.store.repositories.settings import SettingsRepository + + settings = await SettingsRepository(client.store).get_current_settings() + return settings["embeddings"]["model"]["name"] + + +async def test_rebuild_set_embedder_adopts_identity_without_reembedding(temp_db_path): + """SET_EMBEDDER updates stored embedder identity and leaves vectors untouched.""" + from haiku.rag.config import AppConfig + + dim = AppConfig().embeddings.model.vector_dim + sentinel = [0.5] * dim + + async with HaikuRAG(temp_db_path, create=True) as client: + await _add_chunk(client, sentinel) + + drift = AppConfig() + drift.embeddings.model.name = "different-model" + + async with HaikuRAG(temp_db_path, config=drift, skip_validation=True) as client: + async for _ in client.rebuild_database(mode=RebuildMode.SET_EMBEDDER): + pass + + assert await _stored_embedding_name(client) == "different-model" + + rows = (await client.store.chunks_table.query().to_arrow()).to_pylist() + assert len(rows) == 1 + assert rows[0]["vector"] == pytest.approx(sentinel) + + +async def test_rebuild_set_embedder_works_on_empty_database(temp_db_path): + """SET_EMBEDDER reconciles even with no documents (preflight must be bypassed).""" + from haiku.rag.app import HaikuRAGApp + from haiku.rag.config import AppConfig + + async with HaikuRAG(temp_db_path, create=True): + pass + + drift = AppConfig() + drift.embeddings.model.name = "different-model" + + app = HaikuRAGApp(db_path=temp_db_path, config=drift) + await app.rebuild(mode=RebuildMode.SET_EMBEDDER) + + async with HaikuRAG(temp_db_path, config=drift, skip_validation=True) as client: + assert await _stored_embedding_name(client) == "different-model" + + +async def test_rebuild_set_embedder_raises_on_vector_dim_mismatch(temp_db_path): + """SET_EMBEDDER refuses when the vector dimension changed — a full rebuild is needed.""" + from haiku.rag.config import AppConfig + from haiku.rag.store.repositories.settings import ConfigMismatchError + + async with HaikuRAG(temp_db_path, create=True): + pass + + drift = AppConfig() + drift.embeddings.model.vector_dim = 9999 + + async with HaikuRAG(temp_db_path, config=drift, skip_validation=True) as client: + with pytest.raises(ConfigMismatchError): + async for _ in client.rebuild_database(mode=RebuildMode.SET_EMBEDDER): + pass