have rebuild --rechunk re-chunk from the stored docling blob instead of re-converting from the markdown export
This commit is contained in:
parent
7831057339
commit
875c635a77
2 changed files with 33 additions and 8 deletions
|
|
@ -223,28 +223,27 @@ async def _flush_rebuild_batch(
|
|||
async def _rebuild_rechunk(
|
||||
client: "HaikuRAG", documents: list[Document]
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Re-chunk and re-embed from existing document content."""
|
||||
"""Re-chunk and re-embed each document from its stored docling blob."""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
pending_chunks: list[Chunk] = []
|
||||
pending_docs: list[Document] = []
|
||||
pending_doc_ids: list[str] = []
|
||||
|
||||
converter = get_converter(client._config)
|
||||
|
||||
for doc in documents:
|
||||
assert doc.id is not None
|
||||
|
||||
# Convert stored markdown to DoclingDocument
|
||||
docling_document = await converter.convert_text(doc.content, format="md")
|
||||
docling_document = doc.get_docling_document()
|
||||
if docling_document is None:
|
||||
raise ValueError(
|
||||
f"Document {doc.id} has no stored docling document; rechunk "
|
||||
"requires it. Run a full rebuild (without --rechunk) instead."
|
||||
)
|
||||
|
||||
# Chunk and embed
|
||||
chunks = await client.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, client._config)
|
||||
|
||||
# Update document fields
|
||||
doc.set_docling(docling_document)
|
||||
|
||||
# Prepare chunks with document_id and order
|
||||
for order, chunk in enumerate(embedded_chunks):
|
||||
chunk.document_id = doc.id
|
||||
|
|
|
|||
|
|
@ -166,6 +166,32 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat
|
|||
assert "picture" in out.labels
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rechunk_preserves_picture_data(temp_db_path):
|
||||
"""``rebuild --rechunk`` keeps ``picture_data`` for every picture row."""
|
||||
from haiku.rag.client import RebuildMode
|
||||
from haiku.rag.client.documents import _store_document_with_chunks
|
||||
from haiku.rag.store.models.document import Document
|
||||
from tests.store.test_document_items import _docling_doc_with_picture
|
||||
|
||||
docling_doc = _docling_doc_with_picture()
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||
rag._config.processing.pictures = "image"
|
||||
document = Document(content="x", uri="test://doc")
|
||||
document.set_docling(docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
assert created.id is not None
|
||||
before = await rag.document_item_repository.get_all_picture_data(created.id)
|
||||
assert before.get("#/pictures/0") is not None
|
||||
|
||||
async for _ in rag.rebuild_database(mode=RebuildMode.RECHUNK):
|
||||
pass
|
||||
|
||||
after = await rag.document_item_repository.get_all_picture_data(created.id)
|
||||
assert after.get("#/pictures/0") == before.get("#/pictures/0")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_clears_picture_data_when_mode_none(temp_db_path):
|
||||
"""Switching to ``pictures="none"`` and re-running update_document
|
||||
|
|
|
|||
Loading…
Reference in a new issue