Fix embed-only rebuild with changed vector dimensions
When a database was created with one embedding model and rebuild --embed-only was run with a different model, it failed with a vector dimension validation error. - Store reads stored vector_dim when opening existing databases - _rebuild_embed_only recreates chunks table to handle dimension changes - Add test for rebuild with changed vector dimensions
This commit is contained in:
Yiorgis Gozadinos
parent
05be72d773
commit
8bdfbe526c
5 changed files with 347 additions and 39 deletions
|
|
@ -7,6 +7,12 @@
|
|||
- When set, tools emit `{state_key: snapshot}` instead of bare state, enabling state merging when multiple agents share state
|
||||
- Default `None` preserves backwards compatibility (bare state emission)
|
||||
- **Page Image Generation Control**: New `generate_page_images` option in `ConversionOptions` to control PDF page image extraction
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Embed-only rebuild with changed vector dimensions**: Fixed `haiku-rag rebuild --embed-only` failing when the configured embedding model has different dimensions than the database
|
||||
- Store now reads stored vector dimension when opening existing databases, allowing chunks to be read regardless of current config
|
||||
- `_rebuild_embed_only` recreates the chunks table to handle dimension changes
|
||||
- `generate_page_images: bool = True` - Enable/disable rendered page images (used by `visualize_chunk()`)
|
||||
- Works with both `docling-local` and `docling-serve` converters
|
||||
- For `docling-serve`, maps to `image_export_mode` API parameter (`embedded`/`placeholder`)
|
||||
|
|
|
|||
|
|
@ -1415,57 +1415,53 @@ class HaikuRAG:
|
|||
"""Re-embed all chunks without changing chunk boundaries."""
|
||||
from haiku.rag.embeddings import contextualize
|
||||
|
||||
batch_size = 50
|
||||
pending_records: list = []
|
||||
pending_doc_ids: list[str] = []
|
||||
# Collect all chunks with new embeddings
|
||||
all_chunk_data: list[tuple[str, dict]] = []
|
||||
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
|
||||
# Get existing chunks
|
||||
chunks = await self.chunk_repository.get_by_document_id(doc.id)
|
||||
if not chunks:
|
||||
yield doc.id
|
||||
continue
|
||||
|
||||
# Generate new embeddings using contextualize for consistency
|
||||
texts = contextualize(chunks)
|
||||
embeddings = await self.chunk_repository.embedder.embed_documents(texts)
|
||||
|
||||
# Build updated records
|
||||
for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
|
||||
pending_records.append(
|
||||
self.store.ChunkRecord(
|
||||
id=chunk.id, # type: ignore[arg-type]
|
||||
document_id=chunk.document_id, # type: ignore[arg-type]
|
||||
content=chunk.content,
|
||||
content_fts=content_fts,
|
||||
metadata=json.dumps(chunk.metadata),
|
||||
order=chunk.order,
|
||||
vector=embedding,
|
||||
all_chunk_data.append(
|
||||
(
|
||||
doc.id,
|
||||
{
|
||||
"id": chunk.id,
|
||||
"document_id": chunk.document_id,
|
||||
"content": chunk.content,
|
||||
"content_fts": content_fts,
|
||||
"metadata": json.dumps(chunk.metadata),
|
||||
"order": chunk.order,
|
||||
"vector": embedding,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
pending_doc_ids.append(doc.id)
|
||||
# Recreate chunks table (handles dimension changes)
|
||||
self.store.recreate_embeddings_table()
|
||||
|
||||
# Flush batch when size reached
|
||||
if len(pending_doc_ids) >= batch_size:
|
||||
if pending_records:
|
||||
self.store.chunks_table.merge_insert(
|
||||
"id"
|
||||
).when_matched_update_all().execute(pending_records)
|
||||
for doc_id in pending_doc_ids:
|
||||
yield doc_id
|
||||
pending_records = []
|
||||
pending_doc_ids = []
|
||||
# Insert all chunks
|
||||
if all_chunk_data:
|
||||
records = [self.store.ChunkRecord(**data) for _, data in all_chunk_data]
|
||||
self.store.chunks_table.add(records)
|
||||
|
||||
# Flush remaining
|
||||
if pending_records:
|
||||
self.store.chunks_table.merge_insert(
|
||||
"id"
|
||||
).when_matched_update_all().execute(pending_records)
|
||||
for doc_id in pending_doc_ids:
|
||||
yield doc_id
|
||||
# Yield all processed doc IDs
|
||||
yielded_docs: set[str] = set()
|
||||
for doc_id, _ in all_chunk_data:
|
||||
if doc_id not in yielded_docs:
|
||||
yielded_docs.add(doc_id)
|
||||
yield doc_id
|
||||
|
||||
# Yield docs with no chunks
|
||||
for doc in documents:
|
||||
if doc.id and doc.id not in yielded_docs:
|
||||
yield doc.id
|
||||
|
||||
async def _flush_rebuild_batch(
|
||||
self, documents: list[Document], chunks: list[Chunk]
|
||||
|
|
|
|||
|
|
@ -90,12 +90,8 @@ class Store:
|
|||
self._before = before
|
||||
# Time-travel mode is always read-only
|
||||
self._read_only = read_only or (before is not None)
|
||||
self.embedder = get_embedder(config=self._config)
|
||||
self._vacuum_lock = asyncio.Lock()
|
||||
|
||||
# Create the ChunkRecord model with the correct vector dimension
|
||||
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
|
||||
|
||||
# Check if database exists (for local filesystem only)
|
||||
is_new_db = False
|
||||
if not self._has_cloud_config():
|
||||
|
|
@ -113,6 +109,19 @@ class Store:
|
|||
# Connect to LanceDB
|
||||
self.db = self._connect_to_lancedb(db_path)
|
||||
|
||||
# For existing databases, read stored vector dimension to create ChunkRecord
|
||||
# that can read existing chunks. For new databases, use config's dimension.
|
||||
stored_vector_dim = None
|
||||
if not is_new_db:
|
||||
stored_vector_dim = self._get_stored_vector_dim()
|
||||
|
||||
# Create embedder with config's dimension (for generating new embeddings)
|
||||
self.embedder = get_embedder(config=self._config)
|
||||
|
||||
# Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB)
|
||||
chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
|
||||
self.ChunkRecord = create_chunk_model(chunk_vector_dim)
|
||||
|
||||
# Initialize tables (creates them if they don't exist)
|
||||
self._init_tables()
|
||||
|
||||
|
|
@ -137,6 +146,35 @@ class Store:
|
|||
"""Whether the store is in read-only mode."""
|
||||
return self._read_only
|
||||
|
||||
def _get_stored_vector_dim(self) -> int | None:
|
||||
"""Read the stored vector dimension from the settings table.
|
||||
|
||||
Returns:
|
||||
The stored vector dimension, or None if not found.
|
||||
"""
|
||||
try:
|
||||
existing_tables = self.db.table_names()
|
||||
if "settings" not in existing_tables:
|
||||
return None
|
||||
|
||||
settings_table = self.db.open_table("settings")
|
||||
rows = (
|
||||
settings_table.search()
|
||||
.where("id = 'settings'")
|
||||
.limit(1)
|
||||
.to_arrow()
|
||||
.to_pylist()
|
||||
)
|
||||
if not rows or not rows[0].get("settings"):
|
||||
return None
|
||||
|
||||
settings = json.loads(rows[0]["settings"])
|
||||
embeddings = settings.get("embeddings", {})
|
||||
model = embeddings.get("model", {})
|
||||
return model.get("vector_dim")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _assert_writable(self) -> None:
|
||||
"""Raise ReadOnlyError if the store is in read-only mode."""
|
||||
if self._read_only:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -106,6 +106,116 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
|
|||
assert embeddings_before[chunk_id] == embeddings_after[chunk_id]
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_embed_only_with_changed_vector_dim(
|
||||
qa_corpus: Dataset, temp_db_path
|
||||
):
|
||||
"""Test embed-only rebuild when vector dimension changes.
|
||||
|
||||
This tests the scenario where a database was created with one embedding model
|
||||
(e.g., qwen3-embedding:8b with 4096 dims) and rebuild is run with a different
|
||||
model (e.g., qwen3-embedding:4b with 2560 dims).
|
||||
|
||||
The Store should use the stored vector_dim for reading existing chunks,
|
||||
then rebuild should handle changing to the new dimension.
|
||||
"""
|
||||
import json
|
||||
|
||||
import lancedb
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
from pydantic import Field
|
||||
|
||||
# Step 1: Create a database with normal 2560-dim embeddings
|
||||
async with HaikuRAG(temp_db_path, create=True) 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_data = [
|
||||
{
|
||||
"id": c.id,
|
||||
"document_id": c.document_id,
|
||||
"content": c.content,
|
||||
"content_fts": c.content,
|
||||
"metadata": json.dumps(c.metadata),
|
||||
"order": c.order,
|
||||
}
|
||||
for c in chunks_before
|
||||
]
|
||||
|
||||
# Step 2: Manually recreate chunks table with 4096-dim vectors (simulating old DB)
|
||||
db = lancedb.connect(temp_db_path)
|
||||
|
||||
class ChunkRecord4096(LanceModel):
|
||||
id: str
|
||||
document_id: str
|
||||
content: str
|
||||
content_fts: str = Field(default="")
|
||||
metadata: str = Field(default="{}")
|
||||
order: int = Field(default=0)
|
||||
vector: Vector(4096) = Field(default_factory=lambda: [0.0] * 4096) # type: ignore
|
||||
|
||||
db.drop_table("chunks")
|
||||
chunks_table = db.create_table("chunks", schema=ChunkRecord4096)
|
||||
|
||||
# Insert chunks with 4096-dim fake vectors
|
||||
records_4096 = [
|
||||
ChunkRecord4096(
|
||||
id=c["id"],
|
||||
document_id=c["document_id"],
|
||||
content=c["content"],
|
||||
content_fts=c["content_fts"],
|
||||
metadata=c["metadata"],
|
||||
order=c["order"],
|
||||
vector=[0.1] * 4096,
|
||||
)
|
||||
for c in chunk_data
|
||||
]
|
||||
chunks_table.add(records_4096)
|
||||
|
||||
# Update settings to reflect the 4096-dim model used
|
||||
settings_table = db.open_table("settings")
|
||||
rows = (
|
||||
settings_table.search().where("id = 'settings'").limit(1).to_arrow().to_pylist()
|
||||
)
|
||||
settings = json.loads(rows[0]["settings"])
|
||||
settings["embeddings"]["model"]["vector_dim"] = 4096
|
||||
settings["embeddings"]["model"]["name"] = "qwen3-embedding:8b"
|
||||
settings_table.update(
|
||||
where="id = 'settings'", values={"settings": json.dumps(settings)}
|
||||
)
|
||||
|
||||
# Step 3: Open with skip_validation (different config) and run embed-only rebuild
|
||||
# This should work: Store should use stored vector_dim for reading,
|
||||
# then rebuild should migrate to new dimension
|
||||
async with HaikuRAG(temp_db_path, skip_validation=True) as client:
|
||||
processed_ids = [
|
||||
doc_id
|
||||
async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
|
||||
]
|
||||
assert doc.id in processed_ids
|
||||
|
||||
# Verify chunks now have 2560-dim embeddings (from current config's model)
|
||||
chunks_after = await client.chunk_repository.get_by_document_id(doc.id)
|
||||
assert len(chunks_after) == len(chunks_before)
|
||||
|
||||
# Check that embeddings in DB are now 2560-dim
|
||||
raw_chunks = (
|
||||
client.store.chunks_table.search()
|
||||
.where(f"document_id = '{doc.id}'")
|
||||
.to_arrow()
|
||||
.to_pylist()
|
||||
)
|
||||
for raw_chunk in raw_chunks:
|
||||
assert len(raw_chunk["vector"]) == 2560
|
||||
|
||||
# Chunk IDs should be preserved
|
||||
chunk_ids_before = {c.id for c in chunks_before}
|
||||
chunk_ids_after = {c.id for c in chunks_after}
|
||||
assert chunk_ids_before == chunk_ids_after
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
|
||||
"""Test rechunk rebuild: re-chunks from content without accessing source files."""
|
||||
|
|
|
|||
Loading…
Reference in a new issue